server.py

import asyncio
import websockets
import subprocess
import socket
import json
import time
import os
from datetime import datetime

# ─── Helpers ──────────────────────────────────────────────────────────────────

def ts() -> str:
    """Current time prefix for CLI output."""
    return datetime.now().strftime("[%H:%M:%S]")

def get_local_ip() -> str:
    try:
        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 Exception:
        return "127.0.0.1"

def make_response(status: str, output: str) -> str:
    """Build a JSON response string."""
    return json.dumps({
        "status": status,
        "output": output,
        "ts": int(time.time())
    })

# ─── Security ─────────────────────────────────────────────────────────────────

FORBIDDEN = [
    "format", "del /s", "del /f /s", "rd /s", "rmdir /s",
    "rm -rf", "rm -fr", "mkfs", "dd if=",
    ":(){:|:&};:", "shutdown /f",          # fork bomb, forced shutdown
    "reg delete", "bcdedit",               # registry / bootloader
    "cipher /w", "sfc /scannow /offbootdir",
]

COMMAND_TIMEOUT = 10  # seconds

# ─── Special Commands ──────────────────────────────────────────────────────────

START_TIME = time.time()

SPECIAL: dict[str, str] = {
    # Power
    "shutdown":    "shutdown /s /t 5",
    "restart":     "shutdown /r /t 5",
    "sleep":       "rundll32.exe powrprof.dll,SetSuspendState 0,1,0",
    # Volume   (requires nircmd — silently no-ops if not installed)
    "vol_up":      "nircmd changesysvolume 4000",
    "vol_down":    "nircmd changesysvolume -4000",
    # Brightness (nircmd)
    "bright_up":   "nircmd changebrightness 10",
    "bright_down": "nircmd changebrightness -10",
}

# ─── Client Handler ────────────────────────────────────────────────────────────

async def handle_client(websocket):
    client_ip = websocket.remote_address[0]
    print(f"{ts()} [*] New connection from {client_ip}")
    await websocket.send(make_response("ok", "Connected to Laptop Control Server. Ready."))

    try:
        async for raw_message in websocket:
            message = raw_message.strip()
            print(f"{ts()} [>] {client_ip}: {message}")

            # ── Heartbeat ──────────────────────────────────────────────────────
            if message.lower() == "ping":
                uptime = int(time.time() - START_TIME)
                h, rem = divmod(uptime, 3600)
                m, s   = divmod(rem, 60)
                await websocket.send(make_response("ok", f"pong · uptime {h:02d}:{m:02d}:{s:02d}"))
                continue

            # ── Security check ─────────────────────────────────────────────────
            if any(cmd in message.lower() for cmd in FORBIDDEN):
                print(f"{ts()} [!] Forbidden command blocked: {message}")
                await websocket.send(make_response("error", "Forbidden command."))
                continue

            # ── Resolve special aliases ────────────────────────────────────────
            resolved = SPECIAL.get(message.lower(), message)

            # ── Execute ────────────────────────────────────────────────────────
            try:
                proc = await asyncio.create_subprocess_shell(
                    resolved,
                    stdout=asyncio.subprocess.PIPE,
                    stderr=asyncio.subprocess.PIPE,
                )
                try:
                    stdout_b, stderr_b = await asyncio.wait_for(
                        proc.communicate(), timeout=COMMAND_TIMEOUT
                    )
                    stdout = stdout_b.decode(errors="replace").strip()
                    stderr = stderr_b.decode(errors="replace").strip()
                    output = stdout or stderr or f"'{message}' executed (no output)."
                    status = "error" if (not stdout and stderr) else "ok"
                except asyncio.TimeoutError:
                    proc.kill()
                    output = f"Command timed out after {COMMAND_TIMEOUT}s."
                    status = "error"

                print(f"{ts()} [<] Sending response to {client_ip}")
                await websocket.send(make_response(status, output))

            except Exception as exc:
                msg = f"Execution error: {exc}"
                print(f"{ts()} [!] {msg}")
                await websocket.send(make_response("error", msg))

    except websockets.exceptions.ConnectionClosed:
        print(f"{ts()} [-] Connection closed by {client_ip}")
    except Exception as exc:
        print(f"{ts()} [!] Unexpected error ({client_ip}): {exc}")

# ─── Main ──────────────────────────────────────────────────────────────────────

async def main():
    host      = "0.0.0.0"
    port      = 8765
    local_ip  = get_local_ip()

    print("=" * 44)
    print("   Laptop Control System — Server v2")
    print("=" * 44)
    print(f"{ts()} [*] Starting server…")
    print(f"{ts()} [*] Mobile app IP  : {local_ip}")
    print(f"{ts()} [*] Listening on   : ws://{local_ip}:{port}")
    print(f"{ts()} [*] Command timeout: {COMMAND_TIMEOUT}s")
    print(f"{ts()} [*] Responses      : JSON")
    print("=" * 44)

    async with websockets.serve(handle_client, host, port):
        await asyncio.Future()  # run forever

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print(f"\n{ts()} [*] Server stopped.")