This repository has no description
7.5 kB
212 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], timeout: float = 2) -> 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=timeout) # 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]] | None:
75 """Get all Firefox windows and their workspaces from sway tree.
76
77 Returns None if the sway tree could not be queried (so callers can
78 distinguish "no Firefox windows" from "we don't know"). Otherwise
79 returns the list of Firefox windows, possibly empty.
80 """
81 raw = run(["swaymsg", "-t", "get_tree", "-r"], timeout=5)
82 if not raw:
83 return None
84 try:
85 tree: dict[str, Any] = json.loads(raw)
86 except (json.JSONDecodeError, ValueError):
87 return None
88
89 results: list[dict[str, str]] = []
90
91 def traverse(node: dict[str, Any], workspace: str) -> None:
92 node_type = node.get("type", "")
93 if node_type == "workspace":
94 workspace = node.get("name", workspace)
95 app_id = node.get("app_id", "") or ""
96 if "firefox" in app_id.lower() and node_type == "con":
97 results.append(
98 {
99 "workspace": workspace,
100 "windowTitle": node.get("name", ""),
101 }
102 )
103 for child in node.get("nodes", []) + node.get("floating_nodes", []):
104 traverse(child, workspace)
105
106 traverse(tree, "")
107 return results
108
109
110def handle(msg: dict[str, Any]) -> dict[str, Any]:
111 """Handle a single message from the extension."""
112 mid: int = msg.get("id", 0)
113 cmd = msg.get("cmd", "")
114 if cmd == "current_workspace":
115 return {"id": mid, "workspace": run(["workspace-manager", "workspace"])}
116 if cmd == "list_workspaces":
117 try:
118 ws: list[dict[str, Any]] = json.loads(
119 run(["swaymsg", "-t", "get_workspaces", "-r"])
120 )
121 return {"id": mid, "workspaces": [w["name"] for w in ws]}
122 except (json.JSONDecodeError, KeyError):
123 return {"id": mid, "workspaces": []}
124 if cmd == "goto_workspace":
125 run(["goto-workspace", msg.get("name", "")])
126 return {"id": mid, "ok": True}
127 if cmd == "move_to_workspace":
128 run(["swaymsg", "move", "container", "to", "workspace", msg.get("name", "")])
129 return {"id": mid, "ok": True}
130 if cmd == "map_windows":
131 windows = get_firefox_windows()
132 if windows is None:
133 return {"id": mid, "error": "Failed to query sway tree"}
134 return {"id": mid, "windows": windows}
135 if cmd == "save_tabs":
136 workspace = msg.get("workspace", "")
137 tabs = msg.get("tabs", [])
138 resp = daemon_query({"cmd": "save-tabs", "workspace": workspace, "tabs": tabs})
139 return {"id": mid, "ok": resp.get("ok", False)}
140 if cmd == "get_tabs":
141 workspace = msg.get("workspace", "")
142 resp = daemon_query({"cmd": "get-tabs", "workspace": workspace})
143 return {"id": mid, "tabs": resp.get("tabs", [])}
144 return {"id": mid, "error": f"Unknown: {cmd}"}
145
146
147def connect_subscriber() -> socket.socket | None:
148 """Connect to workspace-manager daemon as a subscriber."""
149 try:
150 sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
151 sock.connect(str(SOCKET_PATH))
152 sock.sendall(b"subscribe")
153 sock.setblocking(False)
154 return sock
155 except OSError as e:
156 print(f"Failed to subscribe to workspace-manager: {e}", file=sys.stderr)
157 return None
158
159
160def main() -> None:
161 """Main loop: handle Firefox messages and workspace-manager events."""
162 sub_sock = connect_subscriber()
163 stdin_fd = sys.stdin.buffer.fileno()
164 sub_buf = b""
165
166 while True:
167 read_fds = [stdin_fd]
168 if sub_sock is not None:
169 read_fds.append(sub_sock.fileno())
170
171 try:
172 readable, _, _ = select.select(read_fds, [], [], 10.0)
173 except (ValueError, OSError):
174 break
175
176 if not readable and sub_sock is None:
177 sub_sock = connect_subscriber()
178 continue
179
180 for fd in readable:
181 if fd == stdin_fd:
182 msg = read_message()
183 if msg is None:
184 return
185 write_message(handle(msg))
186 elif sub_sock is not None and fd == sub_sock.fileno():
187 try:
188 data = sub_sock.recv(4096)
189 if not data:
190 sub_sock.close()
191 sub_sock = connect_subscriber()
192 continue
193 sub_buf += data
194 while b"\n" in sub_buf:
195 line, sub_buf = sub_buf.split(b"\n", 1)
196 event: dict[str, Any] = json.loads(line)
197 write_message(
198 {
199 "id": 0,
200 "event": "workspace_change",
201 "workspace": event.get("new", ""),
202 }
203 )
204 except (OSError, json.JSONDecodeError) as e:
205 print(f"Subscriber error: {e}", file=sys.stderr)
206 if sub_sock is not None:
207 sub_sock.close()
208 sub_sock = None
209
210
211if __name__ == "__main__":
212 main()