client.py

#!/usr/bin/env python3
"""
keycoms — client.py
Connects to admin's WebSocket server.
  • Captures keystrokes → sends to admin (admin sees in terminal only)
  • Receives keystrokes from admin → prints in terminal + injects into active
    window with a configurable delay (default 2s per keystroke for debugging)

Usage:
  python3 client.py [ws://host:port]   default: ws://localhost:8765

On Linux run with sudo for /dev/input capture:
  sudo python3 client.py
"""

import asyncio
import json
import sys
import threading

try:
    import keyboard
except ImportError:
    print("[client] Missing: pip install keyboard")
    sys.exit(1)

try:
    from websockets.asyncio.client import connect
    from websockets.exceptions import ConnectionClosed
except ImportError:
    print("[client] Missing: pip install websockets")
    sys.exit(1)

try:
    import inject
except ImportError:
    print("[client] inject.py not found — place it in the same directory.")
    sys.exit(1)

# ── Config ─────────────────────────────────────────────────────────────────────
INJECT_DELAY = 0.0
SERVER_URL   = sys.argv[1] if len(sys.argv) > 1 else "ws://localhost:8765"

# ── State ──────────────────────────────────────────────────────────────────────
state = {
    "ws":      None,
    "loop":    None,
    "running": True,
}

# ── Display helpers ────────────────────────────────────────────────────────────
DISPLAY_MAP = {"\b": "⌫", "\n": "↵\n", "\t": "→"}

def display_char(value: str) -> str:
    return DISPLAY_MAP.get(value, value)


# ── Injection queue (sequential, with delay) ───────────────────────────────────
inject_queue: asyncio.Queue = None

async def injection_worker():
    while state["running"]:
        try:
            value = await asyncio.wait_for(inject_queue.get(), timeout=1.0)
        except asyncio.TimeoutError:
            continue
        print(display_char(value), end="", flush=True)
        inject.key(value)
        inject_queue.task_done()
        await asyncio.sleep(INJECT_DELAY)


# ── Keystroke capture (client → admin) ────────────────────────────────────────
KEY_MAP = {
    "space":     " ",
    "enter":     "\n",
    "backspace": "\b",
    "tab":       "\t",
}

def on_key(event):
    if event.event_type != keyboard.KEY_DOWN:
        return

    name = event.name

    if name == "esc":
        print("\n[client] ESC — disconnecting...")
        state["running"] = False
        ws = state["ws"]
        loop = state["loop"]
        if ws and loop:
            asyncio.run_coroutine_threadsafe(ws.close(), loop)
        return

    if name in KEY_MAP:
        value = KEY_MAP[name]
    elif len(name) == 1:
        value = name
    else:
        return

    ws = state["ws"]
    loop = state["loop"]
    if ws and loop:
        try:
            asyncio.run_coroutine_threadsafe(
                ws.send(json.dumps({"type": "key", "from": "client", "value": value})),
                loop
            )
        except Exception:
            pass


def start_capture():
    keyboard.hook(on_key)
    keyboard.wait()


# ── Message receiver ───────────────────────────────────────────────────────────
async def receiver(ws):
    try:
        async for raw in ws:
            try:
                data = json.loads(raw)
            except json.JSONDecodeError:
                continue

            msg_type = data.get("type")

            if msg_type == "key":
                value = data.get("value", "")
                await inject_queue.put(value)

            elif msg_type == "error":
                print(f"\n[client] Server: {data.get('msg')}")
    except ConnectionClosed:
        pass


# ── Main ───────────────────────────────────────────────────────────────────────
async def main():
    global inject_queue
    inject_queue = asyncio.Queue()

    print("╔═══════════════════════════════╗")
    print("║        keycoms  client        ║")
    print("╚═══════════════════════════════╝")
    print(f"[client] Injection backend : {inject.BACKEND}")
    print(f"[client] Inject delay      : {'none' if INJECT_DELAY == 0 else f'{INJECT_DELAY}s'} per keystroke")
    print(f"[client] Connecting to     : {SERVER_URL}\n")

    try:
        async with connect(SERVER_URL) as ws:
            state["ws"]   = ws
            state["loop"] = asyncio.get_running_loop()

            t = threading.Thread(target=start_capture, daemon=True)
            t.start()

            print("[client] Connected. Type anywhere — your keystrokes go to admin.")
            print("[client] Admin keystrokes appear below with injection:")
            print("[client] ── incoming from admin ──────────────────────────")

            await asyncio.gather(
                receiver(ws),
                injection_worker(),
            )

    except ConnectionRefusedError:
        print(f"[client] Connection refused. Is admin running on {SERVER_URL}?")
    except Exception as e:
        print(f"[client] Disconnected: {e}")

    print("[client] Session ended.")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n[client] Interrupted.")