sender.py

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "rich",
#     "paho-mqtt",
#     "cryptography",
#     "pynput",
#     "zeroconf",
#     "netifaces",
# ]
# ///

import sys
import subprocess
import importlib.util
import time
import threading
import socket
import json
import random
import string
import os

# --- Auto-Install Dependencies ---
REQUIRED_PACKAGES = [
    "rich", "paho-mqtt", "cryptography", "pynput", "zeroconf", "netifaces"
]

def check_and_install_packages():
    missing = []
    for package in REQUIRED_PACKAGES:
        module_name = package.replace("-", "_") if package != "paho-mqtt" else "paho.mqtt"
        if not importlib.util.find_spec(module_name.split('.')[0]):
            missing.append(package)
    
    if missing:
        print(f"Missing packages: {', '.join(missing)}")
        # Check if pip is available
        if not importlib.util.find_spec("pip"):
            print("Error: 'pip' is not available in this environment.")
            print("If you are using 'uv', dependencies should be installed automatically via script metadata.")
            print("Try running: uv run --no-project sender.py")
            sys.exit(1)
            
        print("Attempting to install via pip...")
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install"] + missing)
            print("Installation complete.")
        except subprocess.CalledProcessError as e:
            print(f"Installation failed: {e}")
            sys.exit(1)

check_and_install_packages()

# --- Imports ---
from rich.live import Live
from rich.layout import Layout
from rich.panel import Panel
from rich.console import Console
from rich.table import Table
from rich.text import Text
from rich.align import Align
from rich.prompt import Prompt
import paho.mqtt.client as mqtt
from cryptography.fernet import Fernet
from pynput import keyboard
from zeroconf import ServiceBrowser, Zeroconf

# --- Configuration ---
SERVICE_TYPE = "_qeboard._tcp.local."
DEFAULT_PORT = 5555
MQTT_BROKER = "broker.hivemq.com"
MQTT_PORT = 1883
TOPIC_PREFIX = "qeboard/v1"

# --- Logic ---

class QeBoardSender:
    def __init__(self):
        self.console = Console()
        self.running = True
        self.mode = "SCANNING" # SCANNING, CONNECTED
        self.remote_active = False
        
        # Network
        self.zeroconf = Zeroconf()
        self.mqtt_client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
        self.mqtt_client.on_connect = self.on_mqtt_connect
        self.mqtt_client.on_disconnect = self.on_mqtt_disconnect
        self.mqtt_client.on_message = self.on_mqtt_message
        
        self.found_devices = {} # id -> {name, type, ip, port, last_seen}
        self.selected_device = None
        self.selected_device_id = None
        self.connection_type = None # LAN, WAN
        self.socket = None
        self.fernet = None
        
        # Input
        self.hotkey_listener = None
        self.blocking_listener = None
        self.ctrl_pressed = False
        
        self.logs = []

    def log(self, message):
        timestamp = time.strftime("%H:%M:%S")
        self.logs.append(f"[{timestamp}] {message}")
        if len(self.logs) > 20:
            self.logs.pop(0)

    # --- Discovery ---
    def start_discovery(self):
        # LAN
        self.browser = ServiceBrowser(self.zeroconf, SERVICE_TYPE, self)
        
        # WAN
        try:
            self.mqtt_client.connect(MQTT_BROKER, MQTT_PORT, 10)
            self.mqtt_client.loop_start()
        except Exception as e:
            self.log(f"MQTT Connect Error: {e}")

    def remove_service(self, zeroconf, type, name):
        pass

    def add_service(self, zeroconf, type, name):
        info = zeroconf.get_service_info(type, name)
        if info:
            dev_id = info.properties.get(b'id', b'').decode() or name
            friendly_name = info.properties.get(b'name', b'').decode() or name.split('.')[0]
            addr = socket.inet_ntoa(info.addresses[0])
            self.found_devices[dev_id] = {
                "name": friendly_name,
                "type": "LAN",
                "ip": addr,
                "port": info.port,
                "last_seen": time.time()
            }

    def update_service(self, zeroconf, type, name):
        pass

    def on_mqtt_connect(self, client, userdata, flags, rc, properties):
        client.subscribe(f"{TOPIC_PREFIX}/lobby")
        # Active Scan: Ask receivers to announce themselves
        client.publish(f"{TOPIC_PREFIX}/lobby", json.dumps({"type": "scan"}))

    def on_mqtt_disconnect(self, client, userdata, flags, rc, properties):
        self.log(f"MQTT Disconnected! RC: {rc}")
        if rc != 0:
            self.log("Reconnecting...")
            try:
                client.reconnect()
            except Exception as e:
                self.log(f"Reconnect failed: {e}")

    def on_mqtt_message(self, client, userdata, msg):
        try:
            payload = json.loads(msg.payload.decode())
            dev_id = payload.get('id')
            if dev_id:
                self.found_devices[dev_id] = {
                    "name": payload.get('name'),
                    "type": "WAN",
                    "last_seen": time.time()
                }
        except:
            pass

    # --- Connection ---
    def connect_to(self, device_id, pairing_code):
        device = self.found_devices.get(device_id)
        if not device:
            return "Device not found!"

        try:
            self.fernet = Fernet(pairing_code.encode())
        except Exception as e:
            return "Invalid Pairing Code!"

        self.selected_device = device
        self.selected_device_id = device_id
        
        if device['type'] == 'LAN':
            self.console.print("[yellow]Trying LAN Connection...[/yellow]")
            if self.connect_lan(device['ip'], device['port']):
                return "Success"
            else:
                self.console.print("[red]LAN Failed. Trying WAN...[/red]")
                if self.connect_wan(device_id):
                    return "Success"
                else:
                    return "Connection Failed (LAN & WAN)"
        else:
            if self.connect_wan(device_id):
                return "Success"
            else:
                return "WAN Connection Failed"

    def connect_lan(self, ip, port):
        try:
            self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.socket.settimeout(3) # 3 second timeout
            self.log(f"Connecting to LAN: {ip}:{port}")
            self.socket.connect((ip, port))
            self.socket.settimeout(None) # Reset to blocking
            self.connection_type = "LAN"
            self.mode = "CONNECTED"
            self.start_local_mode()
            return True
        except Exception as e:
            self.log(f"LAN Connection Failed: {e}")
            return False

    def connect_wan(self, device_id):
        # Check if we are connected to MQTT
        if not self.mqtt_client.is_connected():
            self.log("Cannot connect via WAN: MQTT not connected")
            return False
            
        self.connection_type = "WAN"
        self.mode = "CONNECTED"
        self.start_local_mode()
        return True

    def send_payload(self, data):
        try:
            encrypted = self.fernet.encrypt(json.dumps(data).encode())
            
            if self.connection_type == "LAN":
                length = len(encrypted)
                self.socket.sendall(length.to_bytes(4, 'big') + encrypted)
            else:
                # WAN
                if self.selected_device_id:
                    topic = f"{TOPIC_PREFIX}/device/{self.selected_device_id}"
                    self.mqtt_client.publish(topic, encrypted)
                else:
                    self.log("Error: No Device ID selected")
                
        except Exception as e:
            self.log(f"Send Error: {e}")
            self.switch_to_local()

    # --- Input ---
    def start_local_mode(self):
        if self.blocking_listener:
            self.blocking_listener.stop()
            self.blocking_listener = None
            
        self.remote_active = False
        self.hotkey_listener = keyboard.GlobalHotKeys({
            '<ctrl>+<f12>': self.switch_to_remote
        })
        self.hotkey_listener.start()
        self.log("Switched to LOCAL mode")

    def switch_to_remote(self):
        if self.hotkey_listener:
            self.hotkey_listener.stop()
        
        self.remote_active = True
        self.ctrl_pressed = True # Assume from hotkey
        
        self.blocking_listener = keyboard.Listener(
            on_press=self.on_press_remote,
            on_release=self.on_release_remote,
            suppress=True
        )
        self.blocking_listener.start()
        self.log("Switched to REMOTE mode")

    def switch_to_local(self):
        # Called from thread
        threading.Timer(0.1, self.start_local_mode).start()
        return False # Stop listener

    def on_press_remote(self, key):
        try:
            if key == keyboard.Key.ctrl_l or key == keyboard.Key.ctrl_r:
                self.ctrl_pressed = True

            if key == keyboard.Key.f12 and self.ctrl_pressed:
                return self.switch_to_local()

            self.send_key(key, 'press')
        except Exception as e:
            self.log(f"Input Error: {e}")

    def on_release_remote(self, key):
        try:
            if key == keyboard.Key.ctrl_l or key == keyboard.Key.ctrl_r:
                self.ctrl_pressed = False
            self.send_key(key, 'release')
        except Exception as e:
            self.log(f"Input Error: {e}")

    def send_key(self, key, action):
        try:
            key_data = {}
            if isinstance(key, keyboard.Key):
                key_data['special'] = key.name
            elif hasattr(key, 'vk') and key.vk:
                key_data['vk'] = key.vk
            elif hasattr(key, 'char') and key.char:
                key_data['char'] = key.char
            
            self.log(f"Sending {action} {key}") # Verbose logging enabled
            
            self.send_payload({
                'type': 'key_event',
                'action': action,
                'key': key_data
            })
        except Exception as e:
            self.log(f"Send Key Error: {e}")

    # --- UI ---
    def run_tui(self):
        # Try to import msvcrt for Windows non-blocking input
        try:
            import msvcrt
        except ImportError:
            msvcrt = None

        while self.running:
            self.console.clear()
            
            if self.mode == "SCANNING":
                self.show_scanner(msvcrt)
            elif self.mode == "CONNECTED":
                self.show_dashboard()
            
            time.sleep(0.5)

    def show_scanner(self, msvcrt):
        table = Table(title="Available Devices")
        table.add_column("ID", style="cyan")
        table.add_column("Name", style="green")
        table.add_column("Type", style="magenta")
        table.add_column("Last Seen", style="dim")

        active_devices = {}
        now = time.time()
        for dev_id, dev in self.found_devices.items():
            if now - dev['last_seen'] < 60: # 60s timeout
                active_devices[dev_id] = dev
                table.add_row(dev_id, dev['name'], dev['type'], f"{int(now - dev['last_seen'])}s ago")

        self.console.print(Panel(Align.center(Text("QeBoard Sender", style="bold blue"))))
        self.console.print(table)
        self.console.print("\n[bold]Press 'c' to Connect, 'r' to Rescan, or Ctrl+C to Exit[/bold]")
        
        # Non-blocking check
        if msvcrt and msvcrt.kbhit():
            key = msvcrt.getch().decode().lower()
            if key == 'c':
                dev_id = input("\nEnter Device ID: ").strip()
                if dev_id in active_devices:
                    code = input("Enter Pairing Code: ").strip()
                    result = self.connect_to(dev_id, code)
                    if result == "Success":
                        self.console.print("[green]Connected![/green]")
                    else:
                        self.console.print(f"[red]Failed: {result}[/red]")
                        time.sleep(2)
            elif key == 'r':
                self.mqtt_client.publish(f"{TOPIC_PREFIX}/lobby", json.dumps({"type": "scan"}))
                self.console.print("[yellow]Rescanning...[/yellow]")
                time.sleep(1)
        elif not msvcrt:
            # Fallback for non-Windows (blocking)
            print("Non-Windows: Press Enter to refresh list, or type ID to connect.")
            # Use select or just blocking input with timeout?
            # For simplicity in this script, we just block.
            # But this prevents updates.
            # We can use input() with a prompt, but user has to press enter to see updates.
            pass

    def show_dashboard(self):
        status_color = "red" if self.remote_active else "green"
        status_text = "REMOTE (Input Sent)" if self.remote_active else "LOCAL (Normal)"
        
        layout = Layout()
        layout.split(
            Layout(name="header", size=3),
            Layout(name="main", ratio=1),
            Layout(name="logs", size=10)
        )
        
        layout["header"].update(Panel(Align.center(Text("QeBoard Connected", style="bold blue"))))
        
        main_text = Text()
        main_text.append(f"Connected to: {self.selected_device['name']}\n", style="bold")
        main_text.append(f"Type: {self.connection_type}\n")
        main_text.append(f"\nCURRENT MODE:\n{status_text}\n", style=f"bold {status_color} on white")
        main_text.append("\nPress Ctrl+F12 to Toggle", style="dim")
        
        layout["main"].update(Panel(main_text))
        
        log_text = "\n".join(self.logs[-8:])
        layout["logs"].update(Panel(log_text, title="Logs"))
        
        # Use Live for dashboard to avoid flickering
        with Live(layout, refresh_per_second=4, screen=True) as live:
            while self.mode == "CONNECTED" and self.running:
                status_color = "red" if self.remote_active else "green"
                status_text = "REMOTE (Input Sent)" if self.remote_active else "LOCAL (Normal)"
                
                main_text = Text()
                main_text.append(f"Connected to: {self.selected_device['name']}\n", style="bold")
                main_text.append(f"Type: {self.connection_type}\n")
                main_text.append(f"\nCURRENT MODE:\n{status_text}\n", style=f"bold {status_color} on white")
                main_text.append("\nPress Ctrl+F12 to Toggle", style="dim")
                
                layout["main"].update(Panel(main_text))
                layout["logs"].update(Panel("\n".join(self.logs[-8:]), title="Logs"))
                time.sleep(0.25)

    def run(self):
        self.start_discovery()
        try:
            self.run_tui()
        except KeyboardInterrupt:
            pass
        finally:
            self.running = False
            self.mqtt_client.loop_stop()
            self.zeroconf.close()
            if self.hotkey_listener: self.hotkey_listener.stop()
            if self.blocking_listener: self.blocking_listener.stop()

if __name__ == "__main__":
    sender = QeBoardSender()
    sender.run()