#!/usr/bin/env python3
# Fix Chrome CDP download behavior
# Usage: python3 fix-cdp-download.py [port] [download_path]

import os
import sys
import json
import socket
import struct
import base64
import urllib.request

CDP_PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 9222
DOWNLOAD_PATH = sys.argv[2] if len(sys.argv) > 2 else "/home/browser/Downloads"


def cdp_json(path):
    url = f"http://127.0.0.1:{CDP_PORT}{path}"
    with urllib.request.urlopen(url, timeout=5) as resp:
        return json.loads(resp.read())


def ws_send(ws_url, payload_dict):
    from urllib.parse import urlparse
    parsed = urlparse(ws_url)
    host = parsed.hostname
    port = parsed.port or 80
    path = parsed.path

    sock = socket.create_connection((host, port), timeout=5)
    try:
        key = base64.b64encode(os.urandom(16)).decode()
        sock.sendall((
            f"GET {path} HTTP/1.1\r\n"
            f"Host: {host}:{port}\r\n"
            f"Upgrade: websocket\r\n"
            f"Connection: Upgrade\r\n"
            f"Sec-WebSocket-Key: {key}\r\n"
            f"Sec-WebSocket-Version: 13\r\n\r\n"
        ).encode())

        resp = b""
        while b"\r\n\r\n" not in resp:
            chunk = sock.recv(4096)
            if not chunk:
                break
            resp += chunk

        payload = json.dumps(payload_dict).encode()
        mask_key = os.urandom(4)
        frame = bytearray([0x81])
        if len(payload) < 126:
            frame.append(0x80 | len(payload))
        elif len(payload) < 65536:
            frame.append(0x80 | 126)
            frame.extend(struct.pack(">H", len(payload)))
        else:
            frame.append(0x80 | 127)
            frame.extend(struct.pack(">Q", len(payload)))
        frame.extend(mask_key)
        frame.extend(bytes(b ^ mask_key[i % 4] for i, b in enumerate(payload)))
        sock.sendall(frame)

        raw = sock.recv(65536)
        # parse WebSocket frame to get response text
        if len(raw) > 2:
            opcode = raw[0] & 0x0f
            payload_len = raw[1] & 0x7f
            data_start = 2
            if payload_len == 126:
                payload_len = struct.unpack(">H", raw[2:4])[0]
                data_start = 4
            elif payload_len == 127:
                payload_len = struct.unpack(">Q", raw[2:10])[0]
                data_start = 10
            return raw[data_start:data_start + payload_len].decode("utf-8", errors="replace")
        return ""
    finally:
        sock.close()


def main():
    print(f"Fetching WebSocket URL from port {CDP_PORT}...")
    try:
        version = cdp_json("/json/version")
    except Exception as e:
        print(f"Error: could not connect to Chrome on port {CDP_PORT}. Is Chrome running?")
        print(f"Detail: {e}")
        sys.exit(1)

    ws_url = version.get("webSocketDebuggerUrl")
    if not ws_url:
        print("Error: could not get webSocketDebuggerUrl")
        sys.exit(1)

    print(f"WebSocket: {ws_url}")
    print(f"Setting download behavior to allow, path: {DOWNLOAD_PATH}")

    response = ws_send(ws_url, {
        "id": 1,
        "method": "Browser.setDownloadBehavior",
        "params": {
            "behavior": "allow",
            "downloadPath": DOWNLOAD_PATH
        }
    })

    print(f"Response: {response}")

    try:
        result = json.loads(response)
        if "result" in result:
            print("Done.")
        else:
            print(f"Error: {result}")
            sys.exit(1)
    except Exception:
        print(f"Error: could not parse response: {response}")
        sys.exit(1)


if __name__ == "__main__":
    main()