from core.crypto import CryptoContext from core.device import Device from core.registry import DeviceRegistry from log.manager import LogManager from utils.display import Display from proto.c2_pb2 import Command, AgentMessage, AgentMsgType # Forward declaration for type hinting to avoid circular import from typing import TYPE_CHECKING if TYPE_CHECKING: from cli.cli import CLI class Transport: def __init__(self, registry: DeviceRegistry, logger: LogManager, cli_instance: 'CLI' = None): self.crypto = CryptoContext() self.registry = registry self.logger = logger self.cli = cli_instance # CLI instance for callback self.command_responses = {} # To track command responses def set_cli(self, cli_instance: 'CLI'): self.cli = cli_instance # ================================================== # RX (ESP → C2) # ================================================== def handle_incoming(self, sock, addr, raw_data: bytes): """ raw_data = BASE64( ChaCha20( Protobuf AgentMessage ) ) """ # Removed verbose transport debug prints # 1) base64 decode try: cipher = self.crypto.b64_decode(raw_data) except Exception as e: Display.error(f"Base64 decode failed from {addr}: {e}") return # 2) chacha decrypt try: protobuf_bytes = self.crypto.decrypt(cipher) except Exception as e: Display.error(f"Decrypt failed from {addr}: {e}") return # 3) protobuf decode → AgentMessage try: msg = AgentMessage.FromString(protobuf_bytes) except Exception as e: Display.error(f"Protobuf decode failed from {addr}: {e}") return if not msg.device_id: Display.error("AgentMessage received without device_id") return self._dispatch(sock, addr, msg) # ================================================== # DISPATCH # ================================================== def _dispatch(self, sock, addr, msg: AgentMessage): device = self.registry.get(msg.device_id) is_new_device = False if not device: device = Device( id=msg.device_id, sock=sock, address=addr ) self.registry.add(device) Display.device_event(device.id, f"Connected from {addr[0]}") is_new_device = True else: # Device reconnected with new socket - update connection info if device.sock != sock: try: device.sock.close() except Exception: pass device.sock = sock device.address = addr Display.device_event(device.id, f"Reconnected from {addr[0]}:{addr[1]}") device.touch() self._handle_agent_message(device, msg) # Auto-query system_info on new device connection if is_new_device: self._auto_query_system_info(device) def _auto_query_system_info(self, device: Device): """Send system_info command automatically when device connects.""" try: cmd = Command() cmd.device_id = device.id cmd.command_name = "system_info" cmd.request_id = f"auto-sysinfo-{device.id}" self.send_command(device.sock, cmd) except Exception as e: Display.error(f"Auto system_info failed for {device.id}: {e}") def _parse_system_info(self, device: Device, payload: str): """Parse system_info response and update device info.""" # Format: chip=esp32 cores=2 flash=external heap=4310096 uptime=7s modules=network,fakeap try: for part in payload.split(): if "=" in part: key, value = part.split("=", 1) if key == "chip": device.chip = value elif key == "modules": device.modules = value # Notify TUI about device info update Display.device_event(device.id, f"INFO: {payload}") # Send special message to update TUI title from utils.display import Display as Disp if Disp._tui_mode: from tui.bridge import tui_bridge, TUIMessage, MessageType tui_bridge.post_message(TUIMessage( msg_type=MessageType.DEVICE_INFO_UPDATED, device_id=device.id, payload=device.modules )) except Exception as e: Display.error(f"Failed to parse system_info: {e}") # ================================================== # AGENT MESSAGE HANDLER # ================================================== def _handle_agent_message(self, device: Device, msg: AgentMessage): payload_str = "" if msg.payload: try: payload_str = msg.payload.decode(errors="ignore") except Exception: payload_str = repr(msg.payload) if msg.type == AgentMsgType.AGENT_CMD_RESULT: # Check if this is auto system_info response if msg.request_id and msg.request_id.startswith("auto-sysinfo-"): self._parse_system_info(device, payload_str) elif msg.request_id and self.cli: self.cli.handle_command_response(msg.request_id, device.id, payload_str, msg.eof) else: Display.device_event(device.id, f"Command result (no request_id or CLI not set): {payload_str}") elif msg.type == AgentMsgType.AGENT_INFO: # Check for system_info response (format: chip=... modules=...) if "chip=" in payload_str and "modules=" in payload_str: self._parse_system_info(device, payload_str) return # Check for MLAT data (format: MLAT:x;y;rssi) elif payload_str.startswith("MLAT:") and self.cli: mlat_data = payload_str[5:] # Remove "MLAT:" prefix if self.cli.mlat_engine.parse_mlat_message(device.id, mlat_data): # Recalculate position if we have enough scanners state = self.cli.mlat_engine.get_state() if state["scanners_count"] >= 3: self.cli.mlat_engine.calculate_position() else: Display.device_event(device.id, f"MLAT: Invalid data format: {mlat_data}") else: Display.device_event(device.id, f"INFO: {payload_str}") elif msg.type == AgentMsgType.AGENT_ERROR: Display.device_event(device.id, f"ERROR: {payload_str}") elif msg.type == AgentMsgType.AGENT_LOG: Display.device_event(device.id, f"LOG: {payload_str}") elif msg.type == AgentMsgType.AGENT_DATA: Display.device_event(device.id, f"DATA: {payload_str}") else: Display.device_event(device.id, f"UNKNOWN Message Type ({AgentMsgType.Name(msg.type)}): {payload_str}") # ================================================== # TX (C2 → ESP) # ================================================== def send_command(self, sock, cmd: Command): """ Command → Protobuf → ChaCha20 → Base64 → \\n """ try: proto = cmd.SerializeToString() # Removed verbose transport debug prints # Encrypt cipher = self.crypto.encrypt(proto) # Base64 b64 = self.crypto.b64_encode(cipher) sock.sendall(b64 + b"\n") except Exception as e: Display.error(f"Failed to send command to {cmd.device_id}: {e}")