kit-busrouter/hub/server.py

#!/usr/bin/env python3
"""Busfleet Hub — fleet registration + dashboard server.

Listens on :8080. Routes:
  GET  /                             — fleet dashboard (HTML)
  GET  /health                       — health check
  GET  /api/register/<DEVICE_ID>     — assign tunnel port, return JSON
  GET  /api/fleet                    — all registered devices with
                                        tunnel status + access links
  POST /api/authorize-key            — authorize device tunnel key
"""

import json
import os
import socket
import subprocess
import sys
from http.server import HTTPServer, BaseHTTPRequestHandler

REGISTER_SH = "/opt/busfleet-hub/register.sh"
REGISTRY    = "/opt/busfleet-hub/port-registry.json"
DASHBOARD   = "/opt/busfleet-hub/dashboard.html"
HUB_IP      = "162.243.83.36"


# ── Helper: check if a TCP port is open (tunnel active) ──────────
def _port_is_open(port):
    try:
        s = socket.create_connection(("127.0.0.1", port), timeout=2)
        s.close()
        return True
    except Exception:
        return False


class HubHandler(BaseHTTPRequestHandler):
    """Handle hub registration, fleet status, and dashboard requests."""

    def do_GET(self):
        # Fleet dashboard (homepage)
        if self.path in ("/", "/index.html"):
            self.serve_file(DASHBOARD, "text/html; charset=utf-8")
            return

        # Health check
        if self.path == "/health":
            self.send_json(200, '{"status":"ok","service":"busfleet-hub"}')
            return

        # Device registration
        if self.path.startswith("/api/register/"):
            device_id = self.path.split("/api/register/", 1)[1].strip("/")
            if not device_id:
                self.send_json_error(400, "Missing device ID")
                return
            self.call_register_script(device_id)
            return

        # Fleet status — all registered devices
        if self.path == "/api/fleet":
            self.serve_fleet_status()
            return

        self.send_json_error(404, "Not found")

    def do_POST(self):
        # Key authorization
        if self.path.startswith("/api/authorize-key"):
            content_length = int(self.headers.get("Content-Length", 0))
            body = self.rfile.read(content_length).decode() if content_length else "{}"

            try:
                data = json.loads(body)
            except json.JSONDecodeError:
                self.send_json_error(400, "Invalid JSON body")
                return

            device_id = data.get("device_id", "").strip()
            pubkey = data.get("pubkey", "").strip()
            if not device_id or not pubkey:
                self.send_json_error(400, "Missing device_id or pubkey")
                return

            if not pubkey.startswith(("ssh-", "ecdsa-", "sk-")) or len(pubkey) < 80:
                self.send_json_error(400, "Invalid SSH public key format")
                return

            try:
                result = subprocess.run(
                    [REGISTER_SH, "--authorize-key", device_id, pubkey],
                    capture_output=True, text=True, timeout=15,
                )
                status = 200 if result.returncode == 0 else 500
                self.send_json(status, result.stdout.strip())
            except subprocess.TimeoutExpired:
                self.send_json_error(504, "Key authorization timed out")
            except Exception as exc:
                self.send_json_error(500, str(exc))
            return

        self.send_json_error(404, "Not found")

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

    def serve_file(self, path, content_type):
        """Serve a static file."""
        if not os.path.isfile(path):
            self.send_json_error(404, "File not found")
            return
        with open(path, "rb") as f:
            data = f.read()
        self.send_response(200)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(data)))
        self.end_headers()
        self.wfile.write(data)

    def call_register_script(self, device_id):
        """Run register.sh and return its output."""
        try:
            result = subprocess.run(
                [REGISTER_SH, device_id],
                capture_output=True, text=True, timeout=15,
            )
            status = 200 if result.returncode == 0 else 500
            self.send_json(status, result.stdout.strip())
        except subprocess.TimeoutExpired:
            self.send_json_error(504, "Registration timed out")
        except Exception as exc:
            self.send_json_error(500, str(exc))

    def serve_fleet_status(self):
        """Build and return fleet status from port-registry.json +
        live tunnel checks."""
        devices = []
        try:
            with open(REGISTRY) as f:
                registry = json.load(f)
        except Exception:
            registry = {}

        for dev_id, entry in registry.items():
            port = entry.get("tunnel_port", 0)
            tunnel_up = _port_is_open(port) if port else False

            devices.append({
                "device_id": dev_id,
                "tunnel_port": port,
                "assigned": entry.get("assigned", ""),
                "tunnel_up": tunnel_up,
                "ssh_command": f"ssh -p {port} kitadmin@{HUB_IP}",
                "dashboard_url": f"http://{HUB_IP}:{port}" if tunnel_up else None,
            })

        # Sort: tunnel_up first, then by device_id
        devices.sort(key=lambda d: (not d["tunnel_up"], d["device_id"]))

        self.send_json(200, json.dumps({
            "hub": HUB_IP,
            "device_count": len(devices),
            "devices": devices,
        }))

    def send_json(self, status, body):
        """Send a JSON response."""
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Access-Control-Allow-Origin", "*")
        self.end_headers()
        self.wfile.write(body.encode() if isinstance(body, str) else body)

    def send_json_error(self, status, message):
        """Send a JSON error response."""
        self.send_json(status, json.dumps({"error": message}))

    def log_message(self, fmt, *args):
        print(f"[hub] {args[0]}", file=sys.stderr)


def main():
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8080
    server = HTTPServer(("0.0.0.0", port), HubHandler)
    print(f"Busfleet Hub listening on :{port}", flush=True)
    print(f"  Dashboard: http://{HUB_IP}:{port}/", flush=True)
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down.", flush=True)
        server.shutdown()


if __name__ == "__main__":
    main()