""" WebSocket Connection Manager for real-time updates Handles bidirectional communication with standalone web app """ from typing import List, Dict, Any, Optional from fastapi import WebSocket import json import asyncio from pathlib import Path class ConnectionManager: """Manages WebSocket connections for broadcasting logs and status""" def __init__(self): self.active_connections: List[WebSocket] = [] self._loop: asyncio.AbstractEventLoop = None async def connect(self, websocket: WebSocket): await websocket.accept() self.active_connections.append(websocket) def disconnect(self, websocket: WebSocket): if websocket in self.active_connections: self.active_connections.remove(websocket) async def send_personal(self, message: str, websocket: WebSocket): await websocket.send_text(message) async def send_json(self, data: Dict[str, Any], websocket: WebSocket): """Send JSON message to specific client""" await websocket.send_text(json.dumps(data)) async def broadcast(self, message: Dict[str, Any]): """Broadcast message to all connected clients""" if not self.active_connections: return text = json.dumps(message) disconnected = [] for connection in self.active_connections: try: await connection.send_text(text) except Exception: disconnected.append(connection) # Clean up disconnected for conn in disconnected: self.disconnect(conn) async def broadcast_log(self, log: str, level: str = "info"): """Broadcast a log message""" await self.broadcast({ "type": "appendLog", "log": log, "level": level }) async def broadcast_status(self, status: Dict[str, Any]): """Broadcast status update""" await self.broadcast({ "type": "status", **status }) async def broadcast_progress(self, step: int, total: int, name: str, detail: str = ""): """Broadcast progress update""" await self.broadcast({ "type": "progress", "step": step, "totalSteps": total, "stepName": name, "detail": detail }) def sync_broadcast(self, message: Dict[str, Any]): """Synchronous broadcast (for use in non-async code)""" if not self.active_connections: return try: loop = asyncio.get_event_loop() if loop.is_running(): asyncio.create_task(self.broadcast(message)) else: loop.run_until_complete(self.broadcast(message)) except RuntimeError: # No event loop, skip pass def sync_log(self, log: str, level: str = "info"): """Synchronous log broadcast""" self.sync_broadcast({ "type": "appendLog", "log": log, "level": level }) # Global instance manager = ConnectionManager() def get_manager() -> ConnectionManager: """Get the global WebSocket manager""" return manager # === Profile Storage for Standalone === PROFILES_FILE = Path(__file__).parent / "profiles.json" def load_profiles() -> Dict[str, Any]: """Load profiles from JSON file""" if PROFILES_FILE.exists(): try: return json.loads(PROFILES_FILE.read_text(encoding='utf-8')) except Exception: pass return {"profiles": [], "activeProfileId": None} def save_profiles(data: Dict[str, Any]): """Save profiles to JSON file""" PROFILES_FILE.write_text(json.dumps(data, indent=2), encoding='utf-8') def get_active_profile() -> Optional[Dict[str, Any]]: """Get the currently active profile""" data = load_profiles() active_id = data.get("activeProfileId") if not active_id: return None for p in data.get("profiles", []): if p.get("id") == active_id: return p return None # === Command Handler === async def handle_command(command: str, data: Dict[str, Any], websocket: WebSocket): """Handle incoming WebSocket commands from frontend""" ws = get_manager() try: if command == "ping": await ws.send_json({"type": "pong"}, websocket) elif command == "refresh": # Load and send accounts from services.token_service import TokenService token_service = TokenService() tokens = token_service.list_tokens() current = token_service.get_current_token() current_refresh = current.raw_data.get('refreshToken') if current else None accounts = [] for token in tokens: is_active = (current_refresh and token.raw_data.get('refreshToken') == current_refresh) accounts.append({ "filename": token.path.name, "isActive": is_active, "isExpired": token.is_expired, "needsRefresh": token.needs_refresh, "tokenData": { "accountName": token.account_name, "email": token.email, "expiresAt": token.expires_at.isoformat() if token.expires_at else None, "refreshToken": token.raw_data.get('refreshToken') if token.raw_data else None, "clientId": token.raw_data.get('_clientId') if token.raw_data else None, "clientSecret": token.raw_data.get('_clientSecret') if token.raw_data else None }, "usage": { "currentUsage": -1, "usageLimit": 500, "percentageUsed": 0 } }) await ws.send_json({ "type": "accountsLoaded", "accounts": accounts }, websocket) elif command == "switchAccount": email = data.get("email") if email: from services.token_service import TokenService token_service = TokenService() token = token_service.get_token(email) if token: token_service.activate_token(token) await ws.broadcast_log(f"✓ Switched to {token.account_name}") await ws.broadcast({"type": "toast", "message": f"Switched to {token.account_name}", "toastType": "success"}) elif command == "deleteAccount": email = data.get("email") if email: from services.token_service import TokenService token_service = TokenService() token_service.delete_token(email) await ws.broadcast_log(f"🗑 Deleted {email}") elif command == "refreshToken": email = data.get("email") if email: from services.token_service import TokenService token_service = TokenService() token = token_service.get_token(email) if token: token_service.refresh_and_save(token) await ws.broadcast_log(f"🔄 Refreshed {token.account_name}") elif command == "startAutoReg": # Get active profile profile = get_active_profile() if not profile: await ws.broadcast({"type": "toast", "message": "No profile configured", "toastType": "error"}) return # Start autoreg via API from app.api.autoreg import start_autoreg, AutoRegConfig from fastapi import BackgroundTasks imap = profile.get("imap", {}) strategy = profile.get("strategy", {}) config = AutoRegConfig( headless=False, spoofing=True, imapServer=imap.get("server"), imapUser=imap.get("user"), imapPassword=imap.get("password"), emailStrategy=strategy.get("type", "single") ) # Run in background import asyncio from app.api.autoreg import run_autoreg asyncio.create_task(run_autoreg(config)) await ws.broadcast({"type": "updateStatus", "status": json.dumps({"step": 0, "totalSteps": 8, "stepName": "Starting..."})}) elif command == "stopAutoReg": from app.api.autoreg import stop_autoreg await stop_autoreg() elif command == "loadProfiles": data = load_profiles() await ws.send_json({ "type": "profilesLoaded", "profiles": data.get("profiles", []), "activeProfileId": data.get("activeProfileId") }, websocket) elif command == "getActiveProfile": profile = get_active_profile() await ws.send_json({ "type": "activeProfileLoaded", "profile": profile }, websocket) elif command == "setActiveProfile": profile_id = data.get("profileId") profiles_data = load_profiles() profiles_data["activeProfileId"] = profile_id save_profiles(profiles_data) # Find and send the profile profile = None for p in profiles_data.get("profiles", []): if p.get("id") == profile_id: profile = p break await ws.broadcast({ "type": "profilesLoaded", "profiles": profiles_data.get("profiles", []), "activeProfileId": profile_id }) await ws.broadcast({ "type": "activeProfileLoaded", "profile": profile }) await ws.broadcast({"type": "toast", "message": "Profile selected", "toastType": "success"}) elif command == "createProfile": profile = data.get("profile", {}) import uuid profile["id"] = str(uuid.uuid4()) profile["stats"] = {"registered": 0, "failed": 0} profiles_data = load_profiles() profiles_data["profiles"].append(profile) # Auto-select if first profile if len(profiles_data["profiles"]) == 1: profiles_data["activeProfileId"] = profile["id"] save_profiles(profiles_data) await ws.broadcast({ "type": "profilesLoaded", "profiles": profiles_data.get("profiles", []), "activeProfileId": profiles_data.get("activeProfileId") }) await ws.broadcast({"type": "toast", "message": "Profile created", "toastType": "success"}) elif command == "updateProfile": profile = data.get("profile", {}) profile_id = profile.get("id") profiles_data = load_profiles() for i, p in enumerate(profiles_data.get("profiles", [])): if p.get("id") == profile_id: # Preserve stats profile["stats"] = p.get("stats", {"registered": 0, "failed": 0}) profiles_data["profiles"][i] = profile break save_profiles(profiles_data) await ws.broadcast({ "type": "profilesLoaded", "profiles": profiles_data.get("profiles", []), "activeProfileId": profiles_data.get("activeProfileId") }) await ws.broadcast({"type": "toast", "message": "Profile updated", "toastType": "success"}) elif command == "deleteProfile": profile_id = data.get("profileId") profiles_data = load_profiles() profiles_data["profiles"] = [p for p in profiles_data.get("profiles", []) if p.get("id") != profile_id] # Clear active if deleted if profiles_data.get("activeProfileId") == profile_id: profiles_data["activeProfileId"] = None save_profiles(profiles_data) await ws.broadcast({ "type": "profilesLoaded", "profiles": profiles_data.get("profiles", []), "activeProfileId": profiles_data.get("activeProfileId") }) await ws.broadcast({"type": "toast", "message": "Profile deleted", "toastType": "success"}) elif command == "getPatchStatus": # Check if Kiro is patched try: from services.kiro_service import KiroService kiro = KiroService() is_patched = kiro.is_patched() machine_id = kiro.get_machine_id() await ws.send_json({ "type": "patchStatus", "isPatched": is_patched, "currentMachineId": machine_id }, websocket) except Exception as e: await ws.send_json({ "type": "patchStatus", "isPatched": False, "error": str(e) }, websocket) elif command == "getProfile": profile_id = data.get("profileId") profiles_data = load_profiles() profile = None for p in profiles_data.get("profiles", []): if p.get("id") == profile_id: profile = p break await ws.send_json({ "type": "profileLoaded", "profile": profile }, websocket) else: await ws.broadcast_log(f"Unknown command: {command}", "warning") except Exception as e: await ws.broadcast_log(f"Error handling {command}: {str(e)}", "error") import traceback traceback.print_exc()