#!/usr/bin/env python3
"""Exercise the binary SecroAuth transport (no base64)."""
from __future__ import annotations
import base64
import hashlib
import hmac
import json
import os
import socket
import struct
import urllib.request
from pathlib import Path

ROOT = Path("/opt/ws-remote")
SID = "056e573923205845"
HOST, PORT = "127.0.0.1", 8080
HWID = "mock-hwid-0123456789abcdef"
EXECUTOR = "mock-ws-bin"
TOKEN = json.loads((ROOT / "accounts.json").read_text())["users"]["admin"]["token"]
STORE = json.loads((ROOT / "secroauth.json").read_text())
SIGN_KEY = hmac.new(str(STORE["secret"]).encode(), ("socket-v2:" + SID).encode(),
                    hashlib.sha256).digest()


def api(path, payload=None):
    data = None if payload is None else json.dumps(payload).encode()
    req = urllib.request.Request(f"http://{HOST}:{PORT}/api/secroauth/{path}",
                                 data=data, headers={
                                     "Content-Type": "application/json",
                                     "Authorization": "Bearer " + TOKEN})
    with urllib.request.urlopen(req, timeout=30) as resp:
        return json.load(resp)


def sign(text):
    return hmac.new(SIGN_KEY, text.encode(), hashlib.sha256).hexdigest()


class WS:
    def __init__(self, path):
        self.sock = socket.create_connection((HOST, PORT), timeout=15)
        key = base64.b64encode(os.urandom(16)).decode()
        self.sock.sendall((f"GET {path} HTTP/1.1\r\nHost: {HOST}:{PORT}\r\n"
                           "Upgrade: websocket\r\nConnection: Upgrade\r\n"
                           f"Sec-WebSocket-Key: {key}\r\n"
                           "Sec-WebSocket-Version: 13\r\n\r\n").encode())
        buf = b""
        while b"\r\n\r\n" not in buf:
            chunk = self.sock.recv(4096)
            if not chunk:
                raise RuntimeError("closed")
            buf += chunk
        status, rest = buf.split(b"\r\n", 1)
        if b"101" not in status:
            raise RuntimeError(status.decode(errors="replace"))
        self.buf = rest.split(b"\r\n\r\n", 1)[1]

    def _rx(self, n):
        while len(self.buf) < n:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise RuntimeError("closed")
            self.buf += chunk
        out, self.buf = self.buf[:n], self.buf[n:]
        return out

    def recv(self):
        b0, b1 = self._rx(2)
        op = b0 & 0x0F
        ln = b1 & 0x7F
        if ln == 126:
            ln = struct.unpack(">H", self._rx(2))[0]
        elif ln == 127:
            ln = struct.unpack(">Q", self._rx(8))[0]
        data = self._rx(ln) if ln else b""
        if op == 0x8:
            raise RuntimeError("server closed")
        if op == 0x9:
            self.send_raw(b"", 0xA)
            return self.recv()
        if op == 0xA:
            return self.recv()
        return op, data

    def send_raw(self, payload, opcode=0x1):
        mask = os.urandom(4)
        masked = bytes(b ^ mask[i & 3] for i, b in enumerate(payload))
        head = bytearray([0x80 | opcode])
        n = len(payload)
        if n < 126:
            head.append(0x80 | n)
        elif n < 1 << 16:
            head.append(0x80 | 126)
            head += struct.pack(">H", n)
        else:
            head.append(0x80 | 127)
            head += struct.pack(">Q", n)
        self.sock.sendall(bytes(head) + mask + masked)

    def send_json(self, obj):
        self.send_raw(json.dumps(obj).encode(), 0x1)

    def close(self):
        try:
            self.sock.close()
        except OSError:
            pass


def main():
    made = api("keys/create", {"script_id": SID, "count": 1, "days": -1,
                               "hwid_lock": True, "note": "ws bin self-test"})
    key = made["keys"][0]
    print("KEY", key[:8] + "...")
    ws = WS(f"/ws?role=secroauth&script_id={SID}&bin=1")
    try:
        op, data = ws.recv()
        hello = json.loads(data)
        assert op == 0x1 and hello["type"] == "hello", (op, data[:80])
        nonce = hello["nonce"]
        assert hello["sig"] == sign(f"hello:{SID}:{nonce}")
        print("HELLO ok")
        ws.send_json({"type": "auth", "nonce": nonce, "key": key, "hwid": HWID,
                      "executor": EXECUTOR,
                      "sig": sign(f"auth:{SID}:{nonce}:{key}:{HWID}:{EXECUTOR}")})
        # auth_ok is text
        while True:
            op, data = ws.recv()
            if op == 0x2 and data.startswith(b"\x00P"):
                print("BIN_PROBE")
                ws.send_json({"type": "bin_ack"})
                continue
            if op == 0x2 and data.startswith(b"\x00C"):
                # chunks can start arriving immediately
                break
            assert op == 0x1, (op, data[:40])
            msg = json.loads(data)
            if msg.get("type") == "auth_ok":
                xkey = base64.b64decode(msg["key"])
                expected_len = int(msg["len"])
                expected_check = int(msg["check"])
                assert msg["sig"] == sign(
                    f"auth_ok:{SID}:{nonce}:{expected_len}:{expected_check}:{msg['key']}")
                print("AUTH_OK", expected_len)
                continue
            raise RuntimeError(msg)
        chunks = []
        while True:
            if op == 0x2:
                assert data.startswith(b"\x00C"), data[:16]
                chunks.append(data[2:])
            else:
                msg = json.loads(data)
                if msg.get("type") == "done":
                    assert msg["sig"] == sign(f"done:{SID}:{nonce}")
                    print("DONE")
                    break
                raise RuntimeError(msg)
            op, data = ws.recv()
        raw = b"".join(chunks)
        assert len(raw) == expected_len, (len(raw), expected_len)
        source = bytes(b ^ xkey[i % 32] ^ ((i + 1) & 0xFF) ^ (((i + 1) >> 8) & 0xFF)
                       for i, b in enumerate(raw))
        check = 0
        for i in range(0, len(source), 128):
            check = ((check * 33) ^ source[i]) & 0xFFFFFFFF
        assert check == expected_check, (check, expected_check)
        deployed = (ROOT / "secroauth_scripts" / (SID + ".lua")).read_bytes()
        assert source == deployed
        print("ARTIFACT_OK bytes=%d chunks=%d" % (len(source), len(chunks)))
    finally:
        ws.close()
        api("keys/action", {"key": key, "action": "delete"})
        print("KEY deleted")


if __name__ == "__main__":
    main()