receiver.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 receiver.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
import paho.mqtt.client as mqtt
from cryptography.fernet import Fernet
from zeroconf import ServiceInfo, Zeroconf
import netifaces

try:
    from pynput.keyboard import Controller, Key, KeyCode
except ImportError as e:
    print(f"\nError importing pynput: {e}")
    if "X connection" in str(e) or "Display" in str(e):
        print("----------------------------------------------------------------")
        print("Running on Linux without an X server?")
        print("pynput requires a graphical environment (X11) to simulate keys.")
        print("Try setting DISPLAY=:0 if you have a screen attached:")
        print("    export DISPLAY=:0 && python receiver.py")
        print("----------------------------------------------------------------")
    sys.exit(1)

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

# --- Logic ---

class QeBoardReceiver:
    def __init__(self):
        self.keyboard = Controller()
        self.running = True
        self.console = Console()
        self.layout = self.make_layout()
        
        # Identity
        self.device_name = socket.gethostname()
        self.device_id = ''.join(random.choices(string.ascii_uppercase + string.digits, k=6))
        self.pairing_code = Fernet.generate_key().decode() # Use Fernet key as pairing code for simplicity
        self.fernet = Fernet(self.pairing_code.encode())
        
        # 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.logs = []
        self.connected_client = None
        self.mode = "IDLE" # IDLE, CONNECTED

    def make_layout(self):
        layout = Layout()
        layout.split(
            Layout(name="header", size=3),
            Layout(name="main", ratio=1),
            Layout(name="footer", size=3)
        )
        layout["main"].split_row(
            Layout(name="info", ratio=1),
            Layout(name="logs", ratio=2)
        )
        return layout

    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)

    def get_ip_address(self):
        try:
            # Connect to a public DNS server to determine the best outgoing interface
            # We don't actually send data, just creating the socket is enough
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            ip = s.getsockname()[0]
            s.close()
            return ip
        except:
            # Fallback to netifaces if offline
            for interface in netifaces.interfaces():
                try:
                    addrs = netifaces.ifaddresses(interface)
                    if netifaces.AF_INET in addrs:
                        for addr_info in addrs[netifaces.AF_INET]:
                            ip = addr_info['addr']
                            if not ip.startswith('127.'):
                                return ip
                except:
                    pass
            return '127.0.0.1'

    # --- MQTT ---
    def on_mqtt_connect(self, client, userdata, flags, rc, properties):
        self.log(f"Connected to MQTT Broker (RC: {rc})")
        # Subscribe to my specific device topic AND the lobby for scans
        client.subscribe([(f"{TOPIC_PREFIX}/device/{self.device_id}", 0), (f"{TOPIC_PREFIX}/lobby", 0)])
        # Publish presence to lobby
        self.publish_presence()

    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:
            # Handle Lobby Messages (JSON)
            if msg.topic.endswith("/lobby"):
                try:
                    payload = msg.payload.decode()
                    data = json.loads(payload)
                    if data.get('type') == 'scan':
                        self.publish_presence()
                except:
                    pass
                return

            # Handle Device Messages (Encrypted)
            if msg.topic.endswith(self.device_id):
                try:
                    # msg.payload is bytes (Fernet token)
                    decrypted = self.fernet.decrypt(msg.payload)
                    data = json.loads(decrypted.decode())
                    self.process_message(data, "WAN")
                except Exception as e:
                    self.log(f"Decryption failed: {e}")
        except Exception as e:
            self.log(f"MQTT Error: {e}")

    def publish_presence(self):
        if not self.running: return
        payload = json.dumps({
            "name": self.device_name,
            "id": self.device_id,
            "type": "receiver"
        })
        self.mqtt_client.publish(f"{TOPIC_PREFIX}/lobby", payload)
        # Re-publish every 30s
        threading.Timer(30, self.publish_presence).start()

    # --- TCP (LAN) ---
    def start_tcp_server(self):
        server_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        server_sock.bind(('0.0.0.0', DEFAULT_PORT))
        server_sock.listen(1)
        self.log(f"TCP Server listening on {DEFAULT_PORT}")
        
        while self.running:
            try:
                conn, addr = server_sock.accept()
                self.log(f"LAN Connection from {addr}")
                threading.Thread(target=self.handle_tcp_client, args=(conn,)).start()
            except:
                break

    def handle_tcp_client(self, conn):
        try:
            while self.running:
                # Read length prefix
                length_bytes = conn.recv(4)
                if not length_bytes: break
                length = int.from_bytes(length_bytes, 'big')
                data = conn.recv(length)
                if not data: break
                
                # Decrypt
                try:
                    msg = json.loads(self.fernet.decrypt(data).decode())
                    self.process_message(msg, "LAN")
                except Exception as e:
                    self.log(f"LAN Decrypt Error: {e}")
        except Exception as e:
            self.log(f"TCP Error: {e}")
        finally:
            conn.close()

    # --- Processing ---
    def process_message(self, msg, source):
        if msg.get('type') == 'key_event':
            self.mode = f"CONNECTED ({source})"
            self.process_key_event(msg)
        elif msg.get('type') == 'ping':
            pass

    def process_key_event(self, msg):
        key_data = msg.get('key')
        action = msg.get('action')
        
        key = None
        if 'special' in key_data:
            try:
                key = getattr(Key, key_data['special'])
            except AttributeError:
                self.log(f"Unknown special key: {key_data['special']}")
        elif 'vk' in key_data and key_data['vk']:
             key = KeyCode.from_vk(key_data['vk'])
        elif 'char' in key_data and key_data['char']:
            key = key_data['char']

        if key:
            try:
                self.log(f"Simulating {action} {key}")
                if action == 'press':
                    self.keyboard.press(key)
                elif action == 'release':
                    self.keyboard.release(key)
            except Exception as e:
                self.log(f"Key Error: {e}")

    # --- UI ---
    def update_ui(self):
        # Header
        self.layout["header"].update(
            Panel(Align.center(Text("QeBoard Receiver", style="bold magenta")), style="blue")
        )
        
        # Info
        info_text = Text()
        info_text.append(f"Device Name: {self.device_name}\n", style="bold green")
        info_text.append(f"Device ID: {self.device_id}\n", style="bold yellow")
        info_text.append(f"IP Address: {self.get_ip_address()}\n", style="cyan")
        info_text.append(f"\nPAIRING CODE:\n{self.pairing_code}\n", style="bold red on white")
        info_text.append(f"\nStatus: {self.mode}", style="bold white")
        
        self.layout["info"].update(Panel(info_text, title="Connection Info"))
        
        # Logs
        log_text = Text()
        for log in self.logs[-15:]:
            log_text.append(log + "\n")
        self.layout["logs"].update(Panel(log_text, title="Logs"))
        
        # Footer
        self.layout["footer"].update(
            Panel(Align.center("Press Ctrl+C to Exit"), style="dim")
        )

    def run(self):
        # Start Network
        try:
            self.mqtt_client.connect(MQTT_BROKER, MQTT_PORT, 10)
            self.mqtt_client.loop_start()
        except Exception as e:
            self.log(f"Failed to connect to MQTT: {e}")

        # Start TCP
        threading.Thread(target=self.start_tcp_server, daemon=True).start()

        # Start Zeroconf
        try:
            info = ServiceInfo(
                SERVICE_TYPE,
                f"QeBoard-{self.device_id}.{SERVICE_TYPE}",
                addresses=[socket.inet_aton(self.get_ip_address())],
                port=DEFAULT_PORT,
                properties={'id': self.device_id, 'name': self.device_name},
                server=f"{self.device_name}.local.",
            )
            self.zeroconf.register_service(info)
        except Exception as e:
            self.log(f"Zeroconf Error: {e}")

        # UI Loop
        with Live(self.layout, refresh_per_second=4, screen=True) as live:
            try:
                while self.running:
                    self.update_ui()
                    time.sleep(0.25)
            except KeyboardInterrupt:
                pass
        
        # Cleanup
        self.running = False
        self.mqtt_client.loop_stop()
        self.zeroconf.close()
        print("Exited.")

if __name__ == "__main__":
    receiver = QeBoardReceiver()
    receiver.run()