This repository has no description
7.1 kB
201 lines
1#!/usr/bin/env python3
2"""Native messaging host for the Workspace Tabs Firefox extension.
3
4Bridges Firefox extension with workspace-manager daemon and sway.
5Tab persistence goes through the daemon (single writer to workspace files).
6"""
7
8from __future__ import annotations
9
10import json
11import os
12import select
13import socket
14import struct
15import subprocess
16import sys
17from pathlib import Path
18from typing import Any
19
20RUNTIME_DIR = Path(os.environ.get("XDG_RUNTIME_DIR", f"/run/user/{os.getuid()}"))
21SOCKET_PATH = RUNTIME_DIR / "workspace-manager" / "socket"
22
23
24def read_message() -> dict[str, Any] | None:
25 """Read a native messaging message from stdin."""
26 raw = sys.stdin.buffer.read(4)
27 if not raw or len(raw) < 4:
28 return None
29 length = struct.unpack("=I", raw)[0]
30 result: dict[str, Any] = json.loads(sys.stdin.buffer.read(length))
31 return result
32
33
34def write_message(msg: dict[str, Any]) -> None:
35 """Write a native messaging message to stdout."""
36 data = json.dumps(msg).encode()
37 sys.stdout.buffer.write(struct.pack("=I", len(data)) + data)
38 sys.stdout.buffer.flush()
39
40
41def run(cmd: list[str]) -> str:
42 """Run a command and return stripped stdout, or empty string on failure."""
43 try:
44 r = subprocess.run(cmd, capture_output=True, text=True, timeout=2) # noqa: S603
45 except (OSError, subprocess.TimeoutExpired):
46 return ""
47 return r.stdout.strip() if r.returncode == 0 else ""
48
49
50def daemon_query(msg: dict[str, Any]) -> dict[str, Any]:
51 """Send a JSON command to the workspace-manager daemon and return the response."""
52 try:
53 sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
54 sock.connect(str(SOCKET_PATH))
55 sock.settimeout(5.0)
56 sock.sendall(json.dumps(msg).encode())
57 # Shutdown write side so daemon knows we're done sending
58 sock.shutdown(socket.SHUT_WR)
59 data = b""
60 while True:
61 chunk = sock.recv(65536)
62 if not chunk:
63 break
64 data += chunk
65 sock.close()
66 if data:
67 result: dict[str, Any] = json.loads(data.strip())
68 return result
69 except (OSError, json.JSONDecodeError) as e:
70 print(f"Daemon query failed: {e}", file=sys.stderr)
71 return {}
72
73
74def get_firefox_windows() -> list[dict[str, str]]:
75 """Get all Firefox windows and their workspaces from sway tree."""
76 try:
77 tree: dict[str, Any] = json.loads(run(["swaymsg", "-t", "get_tree", "-r"]))
78 except (json.JSONDecodeError, ValueError):
79 return []
80
81 results: list[dict[str, str]] = []
82
83 def traverse(node: dict[str, Any], workspace: str) -> None:
84 node_type = node.get("type", "")
85 if node_type == "workspace":
86 workspace = node.get("name", workspace)
87 app_id = node.get("app_id", "") or ""
88 if "firefox" in app_id.lower() and node_type == "con":
89 results.append(
90 {
91 "workspace": workspace,
92 "windowTitle": node.get("name", ""),
93 }
94 )
95 for child in node.get("nodes", []) + node.get("floating_nodes", []):
96 traverse(child, workspace)
97
98 traverse(tree, "")
99 return results
100
101
102def handle(msg: dict[str, Any]) -> dict[str, Any]:
103 """Handle a single message from the extension."""
104 mid: int = msg.get("id", 0)
105 cmd = msg.get("cmd", "")
106 if cmd == "current_workspace":
107 return {"id": mid, "workspace": run(["workspace-manager", "workspace"])}
108 if cmd == "list_workspaces":
109 try:
110 ws: list[dict[str, Any]] = json.loads(
111 run(["swaymsg", "-t", "get_workspaces", "-r"])
112 )
113 return {"id": mid, "workspaces": [w["name"] for w in ws]}
114 except (json.JSONDecodeError, KeyError):
115 return {"id": mid, "workspaces": []}
116 if cmd == "goto_workspace":
117 run(["goto-workspace", msg.get("name", "")])
118 return {"id": mid, "ok": True}
119 if cmd == "move_to_workspace":
120 run(["swaymsg", "move", "container", "to", "workspace", msg.get("name", "")])
121 return {"id": mid, "ok": True}
122 if cmd == "map_windows":
123 return {"id": mid, "windows": get_firefox_windows()}
124 if cmd == "save_tabs":
125 workspace = msg.get("workspace", "")
126 tabs = msg.get("tabs", [])
127 resp = daemon_query({"cmd": "save-tabs", "workspace": workspace, "tabs": tabs})
128 return {"id": mid, "ok": resp.get("ok", False)}
129 if cmd == "get_tabs":
130 workspace = msg.get("workspace", "")
131 resp = daemon_query({"cmd": "get-tabs", "workspace": workspace})
132 return {"id": mid, "tabs": resp.get("tabs", [])}
133 return {"id": mid, "error": f"Unknown: {cmd}"}
134
135
136def connect_subscriber() -> socket.socket | None:
137 """Connect to workspace-manager daemon as a subscriber."""
138 try:
139 sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
140 sock.connect(str(SOCKET_PATH))
141 sock.sendall(b"subscribe")
142 sock.setblocking(False)
143 return sock
144 except OSError as e:
145 print(f"Failed to subscribe to workspace-manager: {e}", file=sys.stderr)
146 return None
147
148
149def main() -> None:
150 """Main loop: handle Firefox messages and workspace-manager events."""
151 sub_sock = connect_subscriber()
152 stdin_fd = sys.stdin.buffer.fileno()
153 sub_buf = b""
154
155 while True:
156 read_fds = [stdin_fd]
157 if sub_sock is not None:
158 read_fds.append(sub_sock.fileno())
159
160 try:
161 readable, _, _ = select.select(read_fds, [], [], 10.0)
162 except (ValueError, OSError):
163 break
164
165 if not readable and sub_sock is None:
166 sub_sock = connect_subscriber()
167 continue
168
169 for fd in readable:
170 if fd == stdin_fd:
171 msg = read_message()
172 if msg is None:
173 return
174 write_message(handle(msg))
175 elif sub_sock is not None and fd == sub_sock.fileno():
176 try:
177 data = sub_sock.recv(4096)
178 if not data:
179 sub_sock.close()
180 sub_sock = connect_subscriber()
181 continue
182 sub_buf += data
183 while b"\n" in sub_buf:
184 line, sub_buf = sub_buf.split(b"\n", 1)
185 event: dict[str, Any] = json.loads(line)
186 write_message(
187 {
188 "id": 0,
189 "event": "workspace_change",
190 "workspace": event.get("new", ""),
191 }
192 )
193 except (OSError, json.JSONDecodeError) as e:
194 print(f"Subscriber error: {e}", file=sys.stderr)
195 if sub_sock is not None:
196 sub_sock.close()
197 sub_sock = None
198
199
200if __name__ == "__main__":
201 main()