#!/usr/bin/env python3
"""
Pure-Python MITM HTTP/HTTPS proxy for Android.
Uses stdlib + cryptography (no mitmproxy_rs needed).
Intercepts HTTP and HTTPS via CONNECT, logs events to event_server at localhost:8765.
"""
from __future__ import annotations
import asyncio
import datetime
import ipaddress
import json
import os
import ssl
import time
import traceback
import urllib.request
from pathlib import Path

CA_DIR = Path(os.environ.get("MITM_CA_DIR",
              str(Path.home() / ".hermes" / "mitm_ca")))
EVENT_SERVER = os.environ.get("EVENT_SERVER_URL", "http://127.0.0.1:8765/events")
PROXY_PORT = int(os.environ.get("MITM_PORT", "8080"))
INTERCEPT_HOSTS = os.environ.get("MITM_INTERCEPT", "").split(",")


def ensure_ca():
    CA_DIR.mkdir(parents=True, exist_ok=True)
    ca_key_path = CA_DIR / "ca.key"
    ca_cert_path = CA_DIR / "ca.crt"
    if ca_key_path.exists() and ca_cert_path.exists():
        return
    from cryptography import x509
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.primitives.asymmetric import rsa
    from cryptography.x509.oid import NameOID

    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    subject = issuer = x509.Name([
        x509.NameAttribute(NameOID.COMMON_NAME, "Hermes MITM CA"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Hermes Agent"),
    ])
    cert = (
        x509.CertificateBuilder()
        .subject_name(subject).issuer_name(issuer)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(datetime.datetime.utcnow())
        .not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=3650))
        .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
        .add_extension(x509.KeyUsage(
            digital_signature=True, key_cert_sign=True, crl_sign=True,
            content_commitment=False, key_encipherment=False,
            data_encipherment=False, key_agreement=False,
            encipher_only=False, decipher_only=False), critical=True)
        .sign(key, hashes.SHA256())
    )
    ca_key_path.write_bytes(key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.TraditionalOpenSSL,
        serialization.NoEncryption()))
    ca_cert_path.write_bytes(cert.public_bytes(serialization.Encoding.PEM))
    print(f"[mitm] Generated CA: {ca_cert_path}")


def make_host_cert(hostname: str) -> tuple[Path, Path]:
    from cryptography import x509
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.primitives.asymmetric import rsa
    from cryptography.x509.oid import NameOID

    safe = hostname.replace("*", "wildcard").replace(":", "_")
    cert_path = CA_DIR / f"{safe}.crt"
    key_path = CA_DIR / f"{safe}.key"
    if cert_path.exists() and key_path.exists():
        return cert_path, key_path

    ca_key = serialization.load_pem_private_key(
        (CA_DIR / "ca.key").read_bytes(), password=None)
    ca_cert = x509.load_pem_x509_certificate((CA_DIR / "ca.crt").read_bytes())

    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    san_list: list = [x509.DNSName(hostname)]
    try:
        san_list.append(x509.IPAddress(ipaddress.ip_address(hostname)))
    except ValueError:
        pass

    cert = (
        x509.CertificateBuilder()
        .subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, hostname)]))
        .issuer_name(ca_cert.subject)
        .public_key(key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(datetime.datetime.utcnow())
        .not_valid_after(datetime.datetime.utcnow() + datetime.timedelta(days=365))
        .add_extension(x509.SubjectAlternativeName(san_list), critical=False)
        .sign(ca_key, hashes.SHA256())
    )
    key_path.write_bytes(key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.TraditionalOpenSSL,
        serialization.NoEncryption()))
    cert_path.write_bytes(cert.public_bytes(serialization.Encoding.PEM))
    return cert_path, key_path


def post_event(event: dict):
    try:
        payload = json.dumps({
            "type": "network",
            "source": "mitm",
            "app": event.pop("app", None),
            "data": event,
        }).encode()
        req = urllib.request.Request(EVENT_SERVER, data=payload,
                                      headers={"Content-Type": "application/json"})
        urllib.request.urlopen(req, timeout=2)
    except Exception:
        pass


async def pipe(reader: asyncio.StreamReader, writer: asyncio.StreamWriter,
               acc: bytearray | None = None, limit: int = 64 * 1024):
    try:
        while True:
            data = await reader.read(4096)
            if not data:
                break
            writer.write(data)
            await writer.drain()
            if acc is not None and len(acc) < limit:
                acc += data
    except Exception:
        pass
    finally:
        try:
            writer.close()
        except Exception:
            pass


async def handle_connect(client_r: asyncio.StreamReader,
                          client_w: asyncio.StreamWriter,
                          host: str, port: int):
    client_w.write(b"HTTP/1.1 200 Connection established\r\n\r\n")
    await client_w.drain()

    intercept = not INTERCEPT_HOSTS[0] or any(h and h in host for h in INTERCEPT_HOSTS)

    if not intercept:
        try:
            srv_r, srv_w = await asyncio.wait_for(
                asyncio.open_connection(host, port), timeout=10)
        except Exception:
            return
        await asyncio.gather(pipe(client_r, srv_w), pipe(srv_r, client_w),
                             return_exceptions=True)
        return

    try:
        cert_path, key_path = await asyncio.get_event_loop().run_in_executor(
            None, make_host_cert, host)
    except Exception as e:
        print(f"[mitm] cert failed for {host}: {e}")
        try:
            srv_r, srv_w = await asyncio.open_connection(host, port)
            await asyncio.gather(pipe(client_r, srv_w), pipe(srv_r, client_w),
                                 return_exceptions=True)
        except Exception:
            pass
        return

    ssl_srv_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ssl_srv_ctx.load_cert_chain(str(cert_path), str(key_path))
    ssl_srv_ctx.set_ciphers("DEFAULT:@SECLEVEL=1")

    ssl_cli_ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ssl_cli_ctx.check_hostname = False
    ssl_cli_ctx.verify_mode = ssl.CERT_NONE

    transport = client_w.transport
    sock = transport.get_extra_info("socket")
    if sock is None:
        return

    import socket as _socket
    raw_sock = sock.dup()
    transport.close()

    try:
        loop = asyncio.get_event_loop()
        tls_sock = await loop.run_in_executor(
            None, lambda: ssl_srv_ctx.wrap_socket(raw_sock, server_side=True))
        tls_r, tls_w = await asyncio.open_connection(sock=tls_sock)

        srv_r, srv_w = await asyncio.wait_for(
            asyncio.open_connection(host, port, ssl=ssl_cli_ctx), timeout=10)

        req_acc: bytearray = bytearray()
        resp_acc: bytearray = bytearray()
        await asyncio.gather(
            pipe(tls_r, srv_w, req_acc),
            pipe(srv_r, tls_w, resp_acc),
            return_exceptions=True)

        post_event({
            "method": "HTTPS",
            "host": host,
            "port": port,
            "req_len": len(req_acc),
            "resp_len": len(resp_acc),
            "req_preview": req_acc[:512].decode("utf-8", errors="replace"),
            "ts": int(time.time()),
        })
    except Exception as e:
        pass


async def handle_http(client_r: asyncio.StreamReader,
                       client_w: asyncio.StreamWriter,
                       header_buf: bytes, method: str, url: str):
    from urllib.parse import urlparse
    p = urlparse(url)
    host = p.hostname or ""
    port = p.port or 80
    path = (p.path or "/") + (("?" + p.query) if p.query else "")

    # For transparent proxy: URL may be a relative path; extract host from Host header
    if not host:
        for line in header_buf.split(b"\r\n"):
            if line.lower().startswith(b"host:"):
                host_header = line.split(b":", 1)[1].strip().decode("utf-8", errors="replace")
                if ":" in host_header:
                    h, _, p_str = host_header.rpartition(":")
                    host = h
                    try:
                        port = int(p_str)
                    except ValueError:
                        port = 80
                else:
                    host = host_header
                break

    try:
        srv_r, srv_w = await asyncio.wait_for(
            asyncio.open_connection(host, port), timeout=10)
    except Exception:
        client_w.write(b"HTTP/1.1 502 Bad Gateway\r\n\r\n")
        return

    lines = header_buf.split(b"\r\n")
    lines[0] = f"{method} {path} HTTP/1.1".encode()
    lines = [l for l in lines if not l.lower().startswith(b"proxy-")]
    srv_w.write(b"\r\n".join(lines))
    await srv_w.drain()

    req_acc: bytearray = bytearray(header_buf)
    resp_acc: bytearray = bytearray()
    await asyncio.gather(
        pipe(client_r, srv_w, req_acc),
        pipe(srv_r, client_w, resp_acc),
        return_exceptions=True)

    post_event({
        "method": method, "url": url, "host": host, "port": port,
        "req_len": len(req_acc), "resp_len": len(resp_acc),
        "ts": int(time.time()),
    })


async def handle_conn(client_r: asyncio.StreamReader,
                       client_w: asyncio.StreamWriter):
    try:
        first_line = await asyncio.wait_for(client_r.readline(), timeout=15)
        if not first_line:
            return
        buf = bytearray(first_line)
        while True:
            line = await asyncio.wait_for(client_r.readline(), timeout=10)
            buf += line
            if line in (b"\r\n", b"\n", b""):
                break

        parts = first_line.decode("utf-8", errors="replace").strip().split()
        if len(parts) < 3:
            return
        method = parts[0].upper()

        if method == "CONNECT":
            hp = parts[1]
            h, _, p = hp.rpartition(":")
            pnum = int(p) if p.isdigit() else 443
            await handle_connect(client_r, client_w, h, pnum)
        else:
            await handle_http(client_r, client_w, bytes(buf), method, parts[1])
    except (asyncio.TimeoutError, ConnectionResetError, BrokenPipeError):
        pass
    except Exception:
        traceback.print_exc()
    finally:
        try:
            client_w.close()
        except Exception:
            pass


def install_system_ca():
    import subprocess
    ca_cert = CA_DIR / "ca.crt"
    if not ca_cert.exists():
        ensure_ca()
    result = subprocess.run(
        ["openssl", "x509", "-subject_hash_old", "-in", str(ca_cert), "-noout"],
        capture_output=True, text=True)
    h = result.stdout.strip()
    if not h:
        print("[mitm] openssl not found, computing hash manually")
        from cryptography import x509
        from cryptography.hazmat.primitives import hashes
        cert = x509.load_pem_x509_certificate(ca_cert.read_bytes())
        # Use SHA1 of the DER-encoded subject as Android hash (openssl -subject_hash_old)
        import hashlib, struct
        subj_der = cert.subject.public_bytes()
        sha1 = hashlib.sha1(subj_der).digest()
        # Reinterpret as little-endian 32-bit and take lower 8 hex digits
        h_int = struct.unpack("<I", sha1[:4])[0]
        h = f"{h_int:08x}"
    dest = f"/system/etc/security/cacerts/{h}.0"
    try:
        subprocess.run(["mount", "-o", "rw,remount", "/system"], check=True)
        subprocess.run(["cp", str(ca_cert), dest], check=True)
        subprocess.run(["chmod", "644", dest], check=True)
        subprocess.run(["mount", "-o", "ro,remount", "/system"])
        print(f"[mitm] CA installed at {dest}")
        print(f"[mitm] Hash: {h}")
    except Exception as e:
        print(f"[mitm] install_system_ca error: {e}")
        print(f"[mitm] Manual: cp {ca_cert} {dest}")


async def main():
    ensure_ca()
    print(f"[mitm] CA: {CA_DIR/'ca.crt'}")
    print(f"[mitm] Proxy on 0.0.0.0:{PROXY_PORT}")
    server = await asyncio.start_server(handle_conn, "0.0.0.0", PROXY_PORT)
    async with server:
        await server.serve_forever()


if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "install-ca":
        ensure_ca()
        install_system_ca()
    else:
        asyncio.run(main())