"""Cookie sync server for Chrome extension integration."""
import asyncio
import json
import logging
import threading
from dataclasses import dataclass, field
from datetime import datetime
from typing import Optional, Callable
from aiohttp import web
logger = logging.getLogger(__name__)
@dataclass
class CookieStore:
"""Thread-safe cookie store with change callbacks."""
_cookies: dict = field(default_factory=dict)
_lock: threading.Lock = field(default_factory=threading.Lock)
_callbacks: list = field(default_factory=list)
_last_updated: Optional[datetime] = None
def get(self, name: str) -> Optional[str]:
"""Get cookie value by name."""
with self._lock:
return self._cookies.get(name)
def set(self, name: str, value: str) -> None:
"""Set cookie value and notify callbacks."""
with self._lock:
old_value = self._cookies.get(name)
if old_value == value:
return # No change
self._cookies[name] = value
self._last_updated = datetime.now()
logger.info(f"Cookie updated: {name} (length: {len(value)})")
# Notify callbacks outside lock
for callback in self._callbacks:
try:
callback(name, value)
except Exception as e:
logger.error(f"Cookie callback error: {e}")
def get_psid(self) -> Optional[str]:
"""Get __Secure-1PSID cookie."""
return self.get("__Secure-1PSID")
def get_psidts(self) -> Optional[str]:
"""Get __Secure-1PSIDTS cookie."""
return self.get("__Secure-1PSIDTS")
def on_change(self, callback: Callable[[str, str], None]) -> None:
"""Register callback for cookie changes."""
self._callbacks.append(callback)
@property
def last_updated(self) -> Optional[datetime]:
"""Get last update timestamp."""
with self._lock:
return self._last_updated
def to_dict(self) -> dict:
"""Get all cookies as dict."""
with self._lock:
return dict(self._cookies)
# Global cookie store singleton
_cookie_store = CookieStore()
def get_cookie_store() -> CookieStore:
"""Get the global cookie store instance."""
return _cookie_store
class CookieSyncServer:
"""HTTP server for receiving cookies from Chrome extension."""
def __init__(
self,
host: str = "127.0.0.1",
port: int = 8765,
cookie_store: Optional[CookieStore] = None
):
self.host = host
self.port = port
self.cookie_store = cookie_store or get_cookie_store()
self._app: Optional[web.Application] = None
self._runner: Optional[web.AppRunner] = None
self._site: Optional[web.TCPSite] = None
async def _handle_sync(self, request: web.Request) -> web.Response:
"""Handle POST /sync-cookie from Chrome extension."""
try:
data = await request.json()
name = data.get("name")
value = data.get("value")
if not name or not value:
return web.json_response(
{"error": "Missing name or value"},
status=400
)
# Store the cookie
self.cookie_store.set(name, value)
return web.json_response({
"success": True,
"name": name,
"timestamp": datetime.now().isoformat()
})
except json.JSONDecodeError:
return web.json_response(
{"error": "Invalid JSON"},
status=400
)
except Exception as e:
logger.exception("Error handling cookie sync")
return web.json_response(
{"error": str(e)},
status=500
)
async def _handle_status(self, request: web.Request) -> web.Response:
"""Handle GET /status to check server health."""
return web.json_response({
"status": "running",
"cookies": list(self.cookie_store.to_dict().keys()),
"last_updated": (
self.cookie_store.last_updated.isoformat()
if self.cookie_store.last_updated else None
)
})
async def _handle_get_cookies(self, request: web.Request) -> web.Response:
"""Handle GET /cookies to return cookie values."""
cookies = self.cookie_store.to_dict()
return web.json_response({
"success": bool(cookies),
"cookies": cookies
})
async def _handle_cors_preflight(self, request: web.Request) -> web.Response:
"""Handle CORS preflight requests."""
return web.Response(
status=200,
headers={
"Access-Control-Allow-Origin": "*",
"Access-Control-Allow-Methods": "POST, GET, OPTIONS",
"Access-Control-Allow-Headers": "Content-Type",
}
)
@web.middleware
async def _cors_middleware(self, request: web.Request, handler):
"""Add CORS headers to all responses."""
if request.method == "OPTIONS":
return await self._handle_cors_preflight(request)
response = await handler(request)
response.headers["Access-Control-Allow-Origin"] = "*"
return response
def _create_app(self) -> web.Application:
"""Create the aiohttp application."""
app = web.Application(middlewares=[self._cors_middleware])
app.router.add_post("/sync-cookie", self._handle_sync)
app.router.add_get("/status", self._handle_status)
app.router.add_get("/cookies", self._handle_get_cookies)
app.router.add_options("/sync-cookie", self._handle_cors_preflight)
return app
async def start(self) -> None:
"""Start the cookie sync server."""
self._app = self._create_app()
self._runner = web.AppRunner(self._app)
await self._runner.setup()
self._site = web.TCPSite(self._runner, self.host, self.port)
try:
await self._site.start()
logger.info(f"Cookie sync server started at http://{self.host}:{self.port}")
except OSError as e:
if e.errno == 10048 or "address already in use" in str(e).lower():
logger.warning(f"Port {self.port} already in use, cookie sync server disabled")
self._site = None
else:
raise
async def stop(self) -> None:
"""Stop the cookie sync server."""
if self._runner:
await self._runner.cleanup()
logger.info("Cookie sync server stopped")
async def __aenter__(self):
await self.start()
return self
async def __aexit__(self, exc_type, exc_val, exc_tb):
await self.stop()
async def run_cookie_sync_server(host: str = "127.0.0.1", port: int = 8765):
"""Run the cookie sync server standalone."""
server = CookieSyncServer(host=host, port=port)
async with server:
logger.info("Cookie sync server is running. Press Ctrl+C to stop.")
try:
while True:
await asyncio.sleep(3600)
except asyncio.CancelledError:
pass
if __name__ == "__main__":
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
asyncio.run(run_cookie_sync_server())
