import asyncio import logging import os if os.environ.get("CLOUD_ENV", "false").lower() != "true": import wmi import pythoncom else: wmi = None pythoncom = None import re import os import sys import sqlite3 try: import usb.core import usb.util except ImportError: usb = None def get_db_path(): if getattr(sys, 'frozen', False): return os.path.join(os.environ.get('APPDATA', os.path.expanduser('~')), 'JARVIS_OS', 'memory.db') # abspath is backend/services/usb_monitor.py # 1 dirname = backend/services # 2 dirname = backend # 3 dirname = project_root project_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) return os.path.join(project_root, 'memory.db') def init_db(): db_path = get_db_path() try: with sqlite3.connect(db_path) as conn: conn.execute('PRAGMA journal_mode=WAL') conn.execute("CREATE TABLE IF NOT EXISTS usb_allowlist (vendor_id TEXT, product_id TEXT, UNIQUE(vendor_id, product_id))") # Omega Event Bus Table conn.execute("CREATE TABLE IF NOT EXISTS omega_events (id TEXT PRIMARY KEY, domain TEXT, event_type TEXT, description TEXT, persona TEXT, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") conn.execute("CREATE INDEX IF NOT EXISTS idx_omega_events_type ON omega_events(event_type)") # Key Usage Tracking Table conn.execute("CREATE TABLE IF NOT EXISTS key_usage (id TEXT PRIMARY KEY, domain TEXT, tokens_used INTEGER, timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP)") conn.execute("CREATE INDEX IF NOT EXISTS idx_key_usage_domain ON key_usage(domain)") conn.commit() except Exception as e: logging.error(f"Failed to init USB DB: {e}") def get_active_usb_drives(): import os if os.environ.get("CLOUD_ENV", "false").lower() == "true": return [] import wmi import pythoncom pythoncom.CoInitialize() c = wmi.WMI() drives = [] try: for disk in c.Win32_DiskDrive(): if "USB" in disk.InterfaceType: drive_letter = "UNKNOWN" for partition in disk.associators("Win32_DiskDriveToDiskPartition"): for logical in partition.associators("Win32_LogicalDiskToPartition"): drive_letter = logical.Caption vid, pid = extract_vid_pid(disk.PNPDeviceID) drives.append({ "name": disk.Model or disk.Caption, "vid": vid, "pid": pid, "serial": disk.PNPDeviceID, "drive_letter": drive_letter }) except Exception as e: import logging logging.error(f"Error enumerating active USB drives: {e}") return drives def is_allowed(vid, pid): if vid == "UNKNOWN" and pid == "UNKNOWN": return False try: with sqlite3.connect(get_db_path()) as conn: cursor = conn.cursor() cursor.execute("SELECT 1 FROM usb_allowlist WHERE vendor_id = ? AND product_id = ?", (vid, pid)) return cursor.fetchone() is not None except Exception: return False def add_usb_to_allowlist(vid, pid): try: with sqlite3.connect(get_db_path()) as conn: conn.execute("INSERT OR REPLACE INTO usb_allowlist (vendor_id, product_id) VALUES (?, ?)", (vid, pid)) conn.commit() except Exception as e: logging.error(f"Failed to allowlist USB: {e}") def remove_usb_from_allowlist(vid, pid): try: with sqlite3.connect(get_db_path()) as conn: conn.execute("DELETE FROM usb_allowlist WHERE vendor_id = ? AND product_id = ?", (vid, pid)) conn.commit() except Exception as e: logging.error(f"Failed to remove USB from allowlist: {e}") def extract_vid_pid(pnp_id): # Match standard USB\VID_1234&PID_5678 match = re.search(r'VID_([0-9A-Fa-f]{4})&PID_([0-9A-Fa-f]{4})', pnp_id, re.IGNORECASE) if match: return match.group(1).upper(), match.group(2).upper() # Match USBSTOR\Disk&Ven_...&Prod_... match = re.search(r'Ven_([^&]+)&Prod_([^&\\]+)', pnp_id, re.IGNORECASE) if match: return match.group(1).upper(), match.group(2).upper() return "UNKNOWN", "UNKNOWN" async def start_usb_monitor(): import os if os.environ.get("CLOUD_ENV", "false").lower() == "true": return from backend.ws.agent_ws import ws_manager logging.info("Starting Python USB Monitor Loop using PyUSB...") init_db() seen_devices = set() unauthorized_insertions = 0 while True: try: current_devices = set() active_drives_payload = [] def fetch_pyusb(): devices = [] try: for dev in usb.core.find(find_all=True): # Filter by mass storage (bInterfaceClass == 8) is_pendrive = False try: for cfg in dev: for intf in cfg: if intf.bInterfaceClass == 8: is_pendrive = True break if is_pendrive: break except Exception as e: import logging; logging.getLogger(__name__).error(f"Swallowed exception: {e}") if is_pendrive: try: mfg = usb.util.get_string(dev, dev.iManufacturer) if dev.iManufacturer else "Unknown" prod = usb.util.get_string(dev, dev.iProduct) if dev.iProduct else "USB Drive" serial = usb.util.get_string(dev, dev.iSerialNumber) if dev.iSerialNumber else f"{dev.idVendor:04X}:{dev.idProduct:04X}" except Exception: mfg = "Unknown" prod = "USB Drive" serial = f"{dev.idVendor:04X}:{dev.idProduct:04X}" devices.append({ "vid": f"{dev.idVendor:04X}", "pid": f"{dev.idProduct:04X}", "manufacturer": mfg, "product": prod, "serial": serial, "is_pendrive": True }) except Exception as e: logging.error(f"PyUSB enumeration error: {e}") return devices # Run PyUSB enumeration in executor loop = asyncio.get_running_loop() usb_devices = await loop.run_in_executor(None, fetch_pyusb) # Cross-reference with WMI on Windows to get the drive letter (PyUSB cannot natively mount-point) wmi_drive_map = {} if wmi: try: def fetch_wmi_letters(): c = wmi.WMI() mapping = {} for disk in c.Win32_DiskDrive(): if "USB" in disk.InterfaceType: d_letter = "UNKNOWN" for partition in disk.associators("Win32_DiskDriveToDiskPartition"): for logical in partition.associators("Win32_LogicalDiskToPartition"): d_letter = logical.Caption mapping[disk.PNPDeviceID] = d_letter return mapping wmi_drive_map = await loop.run_in_executor(None, fetch_wmi_letters) except Exception as e: import logging; logging.getLogger(__name__).error(f"Swallowed exception: {e}") for dev in usb_devices: vid, pid = dev["vid"], dev["pid"] serial = dev["serial"] current_devices.add(serial) authorized = is_allowed(vid, pid) # Best effort mapping to WMI drive letter if available drive_letter = "E:" # Default fallback for pnp_id, letter in wmi_drive_map.items(): if vid in pnp_id and pid in pnp_id: drive_letter = letter break active_drives_payload.append({ "vendor_id": vid, "product_id": pid, "manufacturer": dev["manufacturer"], "product": dev["product"], "serial": serial, "drive_letter": drive_letter, "is_pendrive": True, "authorized": authorized }) # Emit live device list to the visual panel await ws_manager.broadcast({ "event": "usb:device_list", "payload": {"drives": active_drives_payload} }) new_devices = current_devices - seen_devices for serial in new_devices: dev_match = next((d for d in active_drives_payload if d["serial"] == serial), None) if dev_match and not dev_match["authorized"]: unauthorized_insertions += 1 vid, pid = dev_match["vendor_id"], dev_match["product_id"] logging.warning(f"Unauthorized USB Device Detected: {serial} (VID:{vid} PID:{pid}). Violations: {unauthorized_insertions}") if unauthorized_insertions == 1: await ws_manager.broadcast({"event": "usb:warning", "payload": {"vid": vid, "pid": pid, "device_id": serial}}) elif unauthorized_insertions == 2: ws_manager.restricted_mode = True await ws_manager.broadcast({"event": "usb:restricted", "payload": {"vid": vid, "pid": pid, "device_id": serial}}) else: ws_manager.restricted_mode = True await ws_manager.broadcast({"event": "usb:blocked", "payload": {"vid": vid, "pid": pid, "device_id": serial}}) seen_devices = current_devices except asyncio.CancelledError: logging.info("USB Monitor Loop cancelled, shutting down cleanly.") break except Exception as e: logging.error(f"USB Monitor Error: {e}") await asyncio.sleep(2)