#!/usr/bin/env python3
"""
Minimal HTTP + WebSocket server for the POC.

Serves poc.html on / and handles a WebSocket upgrade on /ws that includes
an HttpOnly Set-Cookie header in the 101 Switching Protocols response.
"""
import base64
import hashlib
import socket
import struct
import threading

HOST, PORT = "127.0.0.1", 8765
WS_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

with open("poc.html", "rb") as f:
    POC_HTML = f.read()


def ws_accept(key: str) -> str:
    digest = hashlib.sha1((key + WS_GUID).encode()).digest()
    return base64.b64encode(digest).decode()


def handle(conn: socket.socket, addr):
    try:
        conn.settimeout(10)
        data = b""
        while b"\r\n\r\n" not in data:
            chunk = conn.recv(4096)
            if not chunk:
                return
            data += chunk
        head = data.split(b"\r\n\r\n", 1)[0].decode("latin1")
        lines = head.split("\r\n")
        request_line = lines[0]
        headers = {}
        for line in lines[1:]:
            if ":" in line:
                k, v = line.split(":", 1)
                headers[k.strip().lower()] = v.strip()

        method, path, _ = request_line.split(" ", 2)

        if path.startswith("/ws") and headers.get("upgrade", "").lower() == "websocket":
            key = headers.get("sec-websocket-key", "")
            accept = ws_accept(key)
            resp = (
                "HTTP/1.1 101 Switching Protocols\r\n"
                "Upgrade: websocket\r\n"
                "Connection: Upgrade\r\n"
                f"Sec-WebSocket-Accept: {accept}\r\n"
                "Set-Cookie: SESSIONID=SECRET-HttpOnly-Value-12345; HttpOnly\r\n"
                "X-Marker: poc-ws\r\n"
                "\r\n"
            )
            conn.sendall(resp.encode())
            # Send a tiny text frame "ok" so the client sees an open socket.
            # FIN=1, opcode=1(text), no mask, len=2, payload="ok"
            conn.sendall(bytes([0x81, 0x02]) + b"ok")
            # Read one client frame (masked) then send close.
            try:
                conn.recv(4096)
            except Exception:
                pass
            # Close frame: FIN=1 opcode=8, len=2, code 1000
            conn.sendall(bytes([0x88, 0x02]) + struct.pack(">H", 1000))
            return

        if path == "/" or path.startswith("/poc.html"):
            body = POC_HTML
            resp = (
                "HTTP/1.1 200 OK\r\n"
                "Content-Type: text/html; charset=utf-8\r\n"
                f"Content-Length: {len(body)}\r\n"
                "Connection: close\r\n"
                "\r\n"
            ).encode() + body
            conn.sendall(resp)
            return

        conn.sendall(b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n")
    except Exception as e:
        print(f"[server] error: {e}")
    finally:
        try:
            conn.close()
        except Exception:
            pass


def main():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind((HOST, PORT))
    s.listen(16)
    print(f"[server] listening on http://{HOST}:{PORT}/")
    while True:
        conn, addr = s.accept()
        threading.Thread(target=handle, args=(conn, addr), daemon=True).start()


if __name__ == "__main__":
    main()