client.py

"""Skyline WebSocket RPC test client.

Usage (as a script):
  python skyline_ws_client.py --url ws://localhost:8080/ws \
    --key de7d6fa0195f8a0d2384da929df069caa1f29673f0361c66561840c9754116e8

Then Ctrl+C to exit.

Usage (as a library):
  from skyline_ws_client import SkylineClient
  client = SkylineClient(url, server_identity_public_key_hex)
  client.connect(start_session_payload={...})
  toggles = client.call("get-feature-toggles", {})
  client.close()
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import signal
import threading
import time
from dataclasses import dataclass
from typing import Any, Dict, Optional

import msgpack
import websockets
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey
from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF
from cryptography.hazmat.primitives.serialization import Encoding, PublicFormat


class ProtocolError(RuntimeError):
    """Raised when Skyline backend returns non-OK status."""

    def __init__(self, status: str, message: str | None = None):
        self.status = status
        self.message = message or ""
        super().__init__(f"Skyline protocol error: {status}: {self.message}")


@dataclass
class StartSessionResult:
    is_activated_now: bool
    activation_date: str
    license_plan: str
    expiry_date: str
    session_token: str


def _mp_pack(obj: Any) -> bytes:
    return msgpack.packb(obj, use_bin_type=True)


def _mp_unpack(blob: bytes) -> Any:
    return msgpack.unpackb(blob, raw=False)


def _handshake_proof(
    client_ecdh_pub: bytes,
    server_ecdh_pub: bytes,
    client_nonce: bytes,
    server_nonce: bytes,
) -> bytes:
    ctx = b"SkylineHandshakeV1"
    return ctx + client_ecdh_pub + server_ecdh_pub + client_nonce + server_nonce


def _derive_session_key(shared_secret: bytes, client_nonce: bytes, server_nonce: bytes) -> bytes:
    salt = client_nonce + server_nonce
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        info=b"SkylineSessionKeyV1",
    )
    return hkdf.derive(shared_secret)


class SkylineClient:
    """Blocking (threaded) Skyline client with a simple .call(...) interface."""

    def __init__(
        self,
        url: str,
        server_identity_public_key_hex: str,
        *,
        handshake_nonce_size: int = 32,
        connect_timeout: float = 10.0,
        ping_interval: float = 20.0,
        ping_timeout: float = 20.0,
    ) -> None:
        self.url = url
        self.server_identity_public_key = bytes.fromhex(server_identity_public_key_hex)
        if len(self.server_identity_public_key) != 32:
            raise ValueError("Server identity public key must be 32 bytes (Ed25519)")

        self.handshake_nonce_size = handshake_nonce_size
        self.connect_timeout = connect_timeout
        self.ping_interval = ping_interval
        self.ping_timeout = ping_timeout

        self._thread: Optional[threading.Thread] = None
        self._loop: Optional[asyncio.AbstractEventLoop] = None

        self._ws: Optional[websockets.WebSocketClientProtocol] = None
        self._session_key: Optional[bytes] = None
        self.session_token: Optional[str] = None

        self._seq: int = 0
        self._pending: dict[int, asyncio.Future] = {}
        self._read_task: Optional[asyncio.Task] = None
        self._write_lock: Optional[asyncio.Lock] = None

        self._closed_evt = threading.Event()

    # ------------------------- public API (blocking) -------------------------

    def connect(self, *, start_session_payload: Optional[Dict[str, Any]] = None) -> StartSessionResult:
        """Connect + handshake + start-session. Returns parsed start-session result."""
        self._ensure_thread()
        fut = asyncio.run_coroutine_threadsafe(
            self._async_connect_and_auth(start_session_payload=start_session_payload),
            self._loop,  # type: ignore[arg-type]
        )
        return fut.result(timeout=self.connect_timeout + 10)

    def call(self, method_name: str, payload: Optional[Dict[str, Any]] = None) -> Any:
        """Call an RPC method and return decoded payload (usually dict / list / None).

        Example:
            toggles = client.call("get-feature-toggles", {})
        """
        if self._loop is None:
            raise RuntimeError("Client is not connected. Call client.connect() first.")
        if payload is None:
            payload = {}
        fut = asyncio.run_coroutine_threadsafe(
            self._rpc_call(method_name=method_name, payload=payload),
            self._loop,
        )
        # Let calls wait longer than connect_timeout; tests might include slow backends.
        return fut.result(timeout=max(self.connect_timeout, 10.0) + 30.0)

    def close(self) -> None:
        """Close websocket and stop background loop."""
        if self._loop is None:
            return
        try:
            fut = asyncio.run_coroutine_threadsafe(self._async_close(), self._loop)
            fut.result(timeout=5)
        finally:
            self._stop_loop_thread()

    def wait_forever(self) -> None:
        """Block until Ctrl+C, keeping the connection open."""
        try:
            while not self._closed_evt.is_set():
                time.sleep(0.5)
        except KeyboardInterrupt:
            pass
        finally:
            self.close()

    def __enter__(self) -> "SkylineClient":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    # ------------------------- internal: thread/loop -------------------------

    def _ensure_thread(self) -> None:
        if self._thread and self._thread.is_alive():
            return

        ready = threading.Event()

        def runner() -> None:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            self._loop = loop
            ready.set()
            try:
                loop.run_forever()
            finally:
                # Best-effort cleanup
                pending = asyncio.all_tasks(loop)
                for t in pending:
                    t.cancel()
                try:
                    loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
                except Exception:
                    pass
                loop.close()

        self._thread = threading.Thread(target=runner, name="SkylineClientLoop", daemon=True)
        self._thread.start()
        ready.wait(timeout=5)
        if self._loop is None:
            raise RuntimeError("Failed to start asyncio loop thread")

    def _stop_loop_thread(self) -> None:
        if self._loop is None:
            return
        try:
            self._loop.call_soon_threadsafe(self._loop.stop)
        except Exception:
            pass
        if self._thread:
            self._thread.join(timeout=2)
        self._loop = None
        self._thread = None

    # ------------------------- internal: connect/auth -------------------------

    async def _async_connect_and_auth(self, *, start_session_payload: Optional[Dict[str, Any]]) -> StartSessionResult:
        if start_session_payload is None:
            # Reasonable defaults for tests
            start_session_payload = {
                "productId": "telepower",
                "productVersion": "7.0.3",
                "language": "en",
                "lastLoginReportedByClient": int(time.time()),
                "deviceSerialNumber": "TEST-DEVICE-001",
                "csProductUUID": "test-uuid-1234",
                "licenseKey": "TEST-LICENSE-KEY",
            }

        # Connect websocket
        self._ws = await websockets.connect(
            self.url,
            ping_interval=self.ping_interval,
            ping_timeout=self.ping_timeout,
            open_timeout=self.connect_timeout,
            max_size=None,  # don't artificially limit message size
        )
        self._write_lock = asyncio.Lock()

        # Handshake (plain msgpack)
        await self._handshake()

        # Start reader loop
        self._read_task = asyncio.create_task(self._read_loop(), name="SkylineReadLoop")

        # start-session must be seq=0 and must be called immediately
        self._seq = 0
        result_payload = await self._rpc_call(
            method_name="start-session",
            payload=start_session_payload,
            session_token=None,
        )

        # Cache token for subsequent calls
        token = None
        if isinstance(result_payload, dict):
            token = result_payload.get("sessionToken")
        if not isinstance(token, str) or not token:
            raise RuntimeError(f"start-session did not return sessionToken, got: {result_payload!r}")

        self.session_token = token

        # Return a typed view (while keeping .call() generic)
        return StartSessionResult(
            is_activated_now=bool(result_payload.get("isActivatedNow")),
            activation_date=str(result_payload.get("activationDate")),
            license_plan=str(result_payload.get("licensePlan")),
            expiry_date=str(result_payload.get("expiryDate")),
            session_token=token,
        )

    async def _async_close(self) -> None:
        self._closed_evt.set()
        # Cancel reader
        if self._read_task:
            self._read_task.cancel()
            try:
                await self._read_task
            except Exception:
                pass
            self._read_task = None

        # Close websocket
        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass
            self._ws = None

        # Fail all pending calls
        for seq, fut in list(self._pending.items()):
            if not fut.done():
                fut.set_exception(ConnectionError("connection closed"))
            self._pending.pop(seq, None)

    async def _handshake(self) -> None:
        assert self._ws is not None

        # X25519 keypair + nonce
        client_priv = X25519PrivateKey.generate()
        client_pub = client_priv.public_key().public_bytes(
            encoding=Encoding.Raw,
            format=PublicFormat.Raw,
        )
        client_nonce = os.urandom(self.handshake_nonce_size)

        # Send plaintext handshake request (msgpack)
        hello = {
            "r": "client-hello",
            "d": {
                "clientEcdhPublic": client_pub,
                "nonce": client_nonce,
            },
        }
        await self._ws.send(_mp_pack(hello))

        # Receive handshake response
        raw = await self._ws.recv()
        if not isinstance(raw, (bytes, bytearray)):
            raise RuntimeError("Handshake failed: server sent non-binary frame")

        resp = _mp_unpack(bytes(raw))
        status = resp.get("S")
        if status != "OK" or resp.get("R") is None:
            raise ProtocolError(status=str(status), message=(resp.get("e") or ""))

        result = resp["R"]
        server_pub = result["serverEcdhPublic"]
        signature = result["signature"]
        server_nonce = result["serverNonce"]

        if not (isinstance(server_pub, (bytes, bytearray)) and isinstance(signature, (bytes, bytearray)) and isinstance(server_nonce, (bytes, bytearray))):
            raise RuntimeError("Handshake failed: unexpected response types")

        # Verify signature
        proof = _handshake_proof(client_pub, bytes(server_pub), client_nonce, bytes(server_nonce))
        try:
            Ed25519PublicKey.from_public_bytes(self.server_identity_public_key).verify(bytes(signature), proof)
        except InvalidSignature as e:
            raise RuntimeError("handshake signature verification failed (possible MITM)") from e

        # Derive session key
        shared = client_priv.exchange(X25519PublicKey.from_public_bytes(bytes(server_pub)))
        self._session_key = _derive_session_key(shared, client_nonce, bytes(server_nonce))

    # ------------------------- internal: rpc -------------------------

    async def _rpc_call(
        self,
        *,
        method_name: str,
        payload: Dict[str, Any],
        session_token: Optional[str] = "__use_default__",
    ) -> Any:
        if self._ws is None or self._session_key is None or self._write_lock is None:
            raise RuntimeError("Client is not connected/handshaked")

        # Decide session token
        token: Optional[str]
        if session_token == "__use_default__":
            token = self.session_token
        else:
            token = session_token

        # Allocate seq + response future
        seq = self._seq
        self._seq += 1
        fut: asyncio.Future = asyncio.get_running_loop().create_future()
        self._pending[seq] = fut

        # Build request packet
        req: Dict[str, Any] = {
            "q": int(seq),
            "t": int(time.time()),
            "r": method_name,
            "d": payload if payload is not None else {},
        }
        if token:
            req["s"] = token

        plain = _mp_pack(req)
        enc = self._encrypt(plain)

        # Send
        async with self._write_lock:
            await self._ws.send(enc)

        # Await response
        try:
            return await fut
        finally:
            self._pending.pop(seq, None)

    def _encrypt(self, plaintext: bytes) -> bytes:
        assert self._session_key is not None
        aesgcm = AESGCM(self._session_key)
        nonce = os.urandom(12)  # Go cipher.NewGCM default nonce size = 12
        ct = aesgcm.encrypt(nonce, plaintext, None)
        return nonce + ct

    def _decrypt(self, encrypted: bytes) -> bytes:
        assert self._session_key is not None
        if len(encrypted) < 12:
            raise ValueError("ciphertext too short")
        nonce, ct = encrypted[:12], encrypted[12:]
        aesgcm = AESGCM(self._session_key)
        return aesgcm.decrypt(nonce, ct, None)

    async def _read_loop(self) -> None:
        assert self._ws is not None
        try:
            async for msg in self._ws:
                if not isinstance(msg, (bytes, bytearray)):
                    continue

                try:
                    plain = self._decrypt(bytes(msg))
                    resp = _mp_unpack(plain)
                except Exception as e:
                    # If we can't decrypt/parse, treat as fatal.
                    for _, fut in list(self._pending.items()):
                        if not fut.done():
                            fut.set_exception(e)
                    break

                seq = resp.get("q")
                status = resp.get("S")
                err_msg = resp.get("e")
                payload = resp.get("R")

                if not isinstance(seq, int):
                    continue

                fut = self._pending.get(seq)
                if fut is None or fut.done():
                    continue

                if status != "OK":
                    fut.set_exception(ProtocolError(status=str(status), message=(err_msg or "")))
                else:
                    fut.set_result(payload)
        finally:
            # Connection is closing; fail pending calls
            for _, fut in list(self._pending.items()):
                if not fut.done():
                    fut.set_exception(ConnectionError("connection closed"))
            self._closed_evt.set()


def _parse_json_arg(s: str) -> Dict[str, Any]:
    try:
        val = json.loads(s)
    except Exception as e:
        raise argparse.ArgumentTypeError(f"Invalid JSON: {e}")
    if not isinstance(val, dict):
        raise argparse.ArgumentTypeError("JSON must be an object (dict)")
    return val


def main() -> None:
    parser = argparse.ArgumentParser(description="Skyline WebSocket RPC test client")
    parser.add_argument(
        "--url",
        default="ws://localhost:8080/ws",
        help="WebSocket URL (e.g. ws://localhost:8080/ws or wss://...)",
    )
    parser.add_argument(
        "--key",
        default="366987a2ebe7576905524ba3024cf813fd10a241bdc65076328bc84504e2bcc1",
        help="Server Ed25519 identity public key (hex, 32 bytes)",
    )
    parser.add_argument(
        "--start-session",
        type=_parse_json_arg,
        default=None,
        help=(
            "Optional JSON object for start-session payload. "
            "Example: --start-session '{" 
            "\"productId\":\"telepower\",\"productVersion\":\"7.0.3\",\"language\":\"en\","
            "\"lastLoginReportedByClient\":1765529203,\"deviceSerialNumber\":\"TEST-DEVICE-001\","
            "\"csProductUUID\":\"test-uuid-1234\",\"licenseKey\":\"TEST-LICENSE-KEY\"}'"
        ),
    )
    parser.add_argument(
        "--probe",
        action="store_true",
        help="After start-session, call get-feature-toggles once and print result.",
    )

    args = parser.parse_args()

    client = SkylineClient(args.url, args.key)
    res = client.connect(start_session_payload=args.start_session)
    print("Connected. start-session OK")
    print("  sessionToken:", res.session_token)
    print("  licensePlan:", res.license_plan)
    print("  expiryDate:", res.expiry_date)

    if args.probe:
        try:
            toggles = client.call("get-feature-toggles", {})
            print("get-feature-toggles =>", toggles)
        except Exception as e:
            print("Probe failed:", e)

    print("\nConnection will stay open. Press Ctrl+C to exit.")
    client.wait_forever()


if __name__ == "__main__":
    # On Windows, Ctrl+C is delivered as KeyboardInterrupt; on POSIX we also trap SIGTERM.
    try:
        signal.signal(signal.SIGTERM, lambda *_: (_ for _ in ()).throw(KeyboardInterrupt()))
    except Exception:
        pass
    main()