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