This repository has no description
14 kB
394 lines
1"""Native messaging bridge that connects browser extension and CLI."""
2
3import asyncio
4import base64
5import json
6import logging
7import struct
8import sys
9from pathlib import Path
10from typing import Any
11
12logger = logging.getLogger(__name__)
13
14
15class NativeMessagingBridge:
16 """Bridge between browser extension (native messaging) and CLI (Unix socket)."""
17
18 def __init__(self) -> None:
19 """Initialize the bridge."""
20 self.cli_clients: set[tuple[asyncio.StreamReader, asyncio.StreamWriter]] = set()
21 self.pending_responses: dict[str, asyncio.Future[Any]] = {}
22 self.message_counter = 0
23 self.stdin_reader: asyncio.StreamReader | None = None
24 self.stdout_writer: asyncio.StreamWriter | None = None
25 self.latest_tab_id: str | None = None # Track the most recently created tab
26
27 async def setup_native_messaging(self) -> None:
28 """Set up stdin/stdout for native messaging."""
29 loop = asyncio.get_event_loop()
30
31 # Create stream reader for stdin
32 self.stdin_reader = asyncio.StreamReader()
33 stdin_protocol = asyncio.StreamReaderProtocol(self.stdin_reader)
34 await loop.connect_read_pipe(lambda: stdin_protocol, sys.stdin)
35
36 # Create stream writer for stdout
37 stdout_transport, stdout_protocol = await loop.connect_write_pipe(
38 asyncio.streams.FlowControlMixin,
39 sys.stdout,
40 )
41 self.stdout_writer = asyncio.StreamWriter(stdout_transport, stdout_protocol, None, loop)
42
43 async def read_native_message(self) -> dict[str, Any] | None:
44 """Read a message from native messaging stdin."""
45 if not self.stdin_reader:
46 return None
47
48 try:
49 # Read message length (4 bytes, little-endian)
50 length_bytes = await self.stdin_reader.readexactly(4)
51 length = struct.unpack("<I", length_bytes)[0]
52
53 # Read message body
54 message_bytes = await self.stdin_reader.readexactly(length)
55 message_data: dict[str, Any] = json.loads(message_bytes.decode("utf-8"))
56 except asyncio.IncompleteReadError:
57 logger.info("Native messaging input closed")
58 return None
59 except Exception:
60 logger.exception("Error reading native message")
61 return None
62 else:
63 return message_data
64
65 async def write_native_message(self, message: dict[str, Any]) -> None:
66 """Write a message to native messaging stdout."""
67 if not self.stdout_writer:
68 return
69
70 try:
71 # Encode message
72 message_bytes = json.dumps(message).encode("utf-8")
73
74 # Write message length (4 bytes, little-endian)
75 self.stdout_writer.write(struct.pack("<I", len(message_bytes)))
76
77 # Write message body
78 self.stdout_writer.write(message_bytes)
79
80 await self.stdout_writer.drain()
81 except Exception:
82 logger.exception("Error writing native message")
83
84 async def process_screenshot_response(
85 self,
86 response: dict[str, Any],
87 output_path: str | None = None,
88 ) -> dict[str, Any]:
89 """Process screenshot response - save to file and return path."""
90 try:
91 # Get data URL from response
92 data_url = response["result"]["screenshot"]
93 if not data_url.startswith("data:image/png;base64,"):
94 return response # Return unchanged if not base64 PNG
95
96 # Extract and decode base64 data
97 base64_data = data_url.split(",")[1]
98 image_data = base64.b64decode(base64_data)
99
100 # Use provided output path or generate default
101 if output_path:
102 screenshot_path = Path(output_path)
103 # Create parent directory if it doesn't exist
104 screenshot_path.parent.mkdir(parents=True, exist_ok=True)
105 else:
106 # Fallback to current directory
107 screenshot_path = Path.cwd() / "screenshot.png"
108
109 # Save screenshot
110 screenshot_path.write_bytes(image_data)
111 logger.info("Saved screenshot to: %s", screenshot_path)
112
113 # Return modified response with file path instead of data
114 return {
115 "id": response.get("id"),
116 "success": True,
117 "result": {
118 "screenshot_path": str(screenshot_path),
119 "message": f"Screenshot saved to {screenshot_path}",
120 },
121 }
122 except Exception as e:
123 logger.exception("Error processing screenshot")
124 return {
125 "id": response.get("id"),
126 "success": False,
127 "error": f"Failed to save screenshot: {e!s}",
128 }
129
130 async def handle_extension_message(self, message: dict[str, Any]) -> None:
131 """Process messages from browser extension."""
132 msg_id = message.get("id")
133 command = message.get("command")
134
135 # Handle commands from extension (e.g., save-screenshot from content script)
136 if command == "save-screenshot":
137 await self._handle_save_screenshot(message)
138 return
139
140 if msg_id and msg_id in self.pending_responses:
141 # This is a response to a CLI request
142 future = self.pending_responses.pop(msg_id)
143
144 # Track the latest tab ID if this was a new-tab command
145 if message.get("success") and "tabId" in message.get("result", {}):
146 self.latest_tab_id = message["result"]["tabId"]
147
148 future.set_result(message)
149 else:
150 # This is a command from the extension - shouldn't happen in our architecture
151 logger.warning("Unexpected command from extension: %s", message)
152
153 async def _handle_save_screenshot(self, message: dict[str, Any]) -> None:
154 """Handle save-screenshot command from extension."""
155 msg_id = message.get("id")
156 params = message.get("params", {})
157 data_url = params.get("screenshot", "")
158 output_path = params.get("output_path")
159
160 if not data_url.startswith("data:image/png;base64,"):
161 await self.write_native_message(
162 {
163 "id": msg_id,
164 "success": False,
165 "error": "Invalid screenshot data URL",
166 },
167 )
168 return
169
170 try:
171 # Extract and decode base64 data
172 base64_data = data_url.split(",")[1]
173 image_data = base64.b64decode(base64_data)
174
175 # Use provided output path or generate default
176 if output_path:
177 screenshot_path = Path(output_path)
178 screenshot_path.parent.mkdir(parents=True, exist_ok=True)
179 else:
180 screenshot_path = Path.cwd() / "screenshot.png"
181
182 # Save screenshot
183 screenshot_path.write_bytes(image_data)
184 logger.info("Saved screenshot to: %s", screenshot_path)
185
186 # Send success response back to extension
187 await self.write_native_message(
188 {
189 "id": msg_id,
190 "success": True,
191 "result": {
192 "screenshot_path": str(screenshot_path),
193 "message": f"Screenshot saved to {screenshot_path}",
194 },
195 },
196 )
197 except Exception as e:
198 logger.exception("Error saving screenshot")
199 await self.write_native_message(
200 {
201 "id": msg_id,
202 "success": False,
203 "error": f"Failed to save screenshot: {e!s}",
204 },
205 )
206
207 async def handle_cli_client(
208 self,
209 reader: asyncio.StreamReader,
210 writer: asyncio.StreamWriter,
211 ) -> None:
212 """Handle connection from CLI client."""
213 logger.info("CLI client connected")
214 client = (reader, writer)
215 self.cli_clients.add(client)
216
217 try:
218 while True:
219 # Read JSON line
220 line = await reader.readline()
221 if not line:
222 break
223
224 await self.handle_cli_message(writer, line.decode("utf-8").strip())
225
226 except Exception:
227 logger.exception("Error handling CLI client")
228 finally:
229 logger.info("CLI client disconnected")
230 self.cli_clients.discard(client)
231 writer.close()
232 await writer.wait_closed()
233
234 async def _ensure_tab_exists(
235 self,
236 data: dict[str, Any],
237 msg_id: str,
238 writer: asyncio.StreamWriter,
239 ) -> bool:
240 """Ensure a tab exists for commands that need one.
241
242 Returns True if tab exists or was created, False if failed.
243 """
244 command = data.get("command", "")
245 if command in ["list-tabs", "new-tab"] or "tabId" in data:
246 return True
247
248 if self.latest_tab_id:
249 data["tabId"] = self.latest_tab_id
250 return True
251
252 # No tabs exist, create one
253 logger.info("No managed tabs exist, creating a new tab")
254 new_tab_id = f"auto_{msg_id}"
255 new_tab_msg = {
256 "command": "new-tab",
257 "params": {"url": "https://example.com"},
258 "id": new_tab_id,
259 }
260 new_tab_future: asyncio.Future[Any] = asyncio.Future()
261 self.pending_responses[new_tab_id] = new_tab_future
262 await self.write_native_message(new_tab_msg)
263
264 try:
265 new_tab_response = await asyncio.wait_for(new_tab_future, timeout=5.0)
266 result = new_tab_response.get("result", {})
267 if new_tab_response.get("success") and "tabId" in result:
268 self.latest_tab_id = new_tab_response["result"]["tabId"]
269 data["tabId"] = self.latest_tab_id
270 logger.info("Created new tab: %s", self.latest_tab_id)
271 return True
272 logger.error("Failed to create new tab - no tabId in response")
273 self.latest_tab_id = None
274 except TimeoutError:
275 logger.exception("Timeout creating new tab")
276 await self._send_error_response(
277 writer,
278 msg_id,
279 "No managed tabs exist. Use 'browser-cli new-tab'",
280 )
281 return False
282 else:
283 return False
284
285 async def _send_error_response(
286 self,
287 writer: asyncio.StreamWriter,
288 msg_id: str,
289 error: str,
290 ) -> None:
291 """Send error response to CLI client."""
292 error_response = {"error": error, "id": msg_id}
293 writer.write((json.dumps(error_response) + "\n").encode("utf-8"))
294 await writer.drain()
295
296 async def _send_response(self, writer: asyncio.StreamWriter, response: dict[str, Any]) -> None:
297 """Send response to CLI client."""
298 writer.write((json.dumps(response) + "\n").encode("utf-8"))
299 await writer.drain()
300
301 async def handle_cli_message(
302 self,
303 writer: asyncio.StreamWriter,
304 message: str,
305 ) -> None:
306 """Process messages from CLI client."""
307 try:
308 data = json.loads(message)
309
310 # Use the client's ID or generate one if missing
311 msg_id = data.get("id")
312 if not msg_id:
313 msg_id = f"cli_{self.message_counter}"
314 self.message_counter += 1
315 data["id"] = msg_id
316
317 # Ensure tab exists for commands that need it
318 if not await self._ensure_tab_exists(data, msg_id, writer):
319 return
320
321 # Create future for response
322 future: asyncio.Future[Any] = asyncio.Future()
323 self.pending_responses[msg_id] = future
324
325 # Forward to extension via native messaging
326 await self.write_native_message(data)
327
328 # Wait for response with timeout
329 try:
330 response = await asyncio.wait_for(future, timeout=30.0)
331
332 # Special handling for screenshot command
333 command = data.get("command", "")
334 if (
335 command == "screenshot"
336 and response.get("success")
337 and response.get("result", {}).get("screenshot")
338 ):
339 output_path = data.get("params", {}).get("output_path")
340 response = await self.process_screenshot_response(response, output_path)
341
342 await self._send_response(writer, response)
343 except TimeoutError:
344 self.pending_responses.pop(msg_id, None)
345 await self._send_error_response(writer, msg_id, "Request timeout")
346
347 except json.JSONDecodeError:
348 logger.exception("Invalid JSON from CLI: %s", message)
349 await self._send_response(writer, {"error": "Invalid JSON"})
350
351 async def native_messaging_loop(self) -> None:
352 """Handle native messaging in a loop."""
353 await self.setup_native_messaging()
354
355 while True:
356 message = await self.read_native_message()
357 if message is None:
358 break
359
360 await self.handle_extension_message(message)
361
362 async def start(self, socket_path: Path) -> None:
363 """Start the bridge server."""
364 # Remove old socket if exists
365 socket_path.unlink(missing_ok=True)
366
367 # Start Unix socket server for CLI
368 server = await asyncio.start_unix_server(
369 self.handle_cli_client,
370 socket_path,
371 )
372
373 # Set socket permissions
374 socket_path.chmod(0o600)
375
376 # Send socket path to extension via native messaging
377 await self.write_native_message(
378 {
379 "socket_path": str(socket_path),
380 "ready": True,
381 },
382 )
383
384 logger.info("Native messaging bridge started")
385 logger.info("CLI socket: %s", socket_path)
386
387 # Run native messaging loop
388 try:
389 await self.native_messaging_loop()
390 finally:
391 # Clean shutdown
392 server.close()
393 await server.wait_closed()
394 socket_path.unlink(missing_ok=True)