#!/usr/bin/env python3
"""Original LockedIn FDE exercise: actual loopback TCP, TLS and HTTP.

No external target, provider key or model. Never install the temporary CA into
your machine's trust store. This small fixture is not a production web server.
"""
from __future__ import annotations

import argparse
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
import ipaddress
import json
import os
from pathlib import Path
import platform
import socket
import ssl
import tempfile
import threading
import unittest

import cryptography
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.x509.oid import ExtendedKeyUsageOID, NameOID

CASES = ("success", "no-listener", "wrong-name", "untrusted", "denied")
LAYERS = ("lookup", "tcp", "tls", "http")
TIMEOUT = 3


def private_file(path: Path, data: bytes) -> None:
    # Exclusive creation inside an owned temporary directory, never a user file.
    with os.fdopen(os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600), "wb") as f:
        f.write(data)


def certificates(folder: Path) -> tuple[Path, Path, Path]:
    """Fresh local authority and localhost leaf; authority private key stays in RAM."""
    now = datetime.now(timezone.utc)
    authority_key = ec.generate_private_key(ec.SECP256R1())
    leaf_key = ec.generate_private_key(ec.SECP256R1())
    authority_name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "Ephemeral FDE lab CA")])
    subject = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "localhost")])

    def builder(name, issuer, key):
        return (x509.CertificateBuilder().subject_name(name).issuer_name(issuer)
                .public_key(key.public_key()).serial_number(x509.random_serial_number())
                .not_valid_before(now - timedelta(minutes=1))
                .not_valid_after(now + timedelta(hours=1)))

    authority = (builder(authority_name, authority_name, authority_key)
                 .add_extension(x509.BasicConstraints(ca=True, path_length=0), critical=True)
                 .add_extension(x509.KeyUsage(True, False, False, False, False, True, True, False, False), critical=True)
                 .add_extension(x509.SubjectKeyIdentifier.from_public_key(authority_key.public_key()), critical=False)
                 .sign(authority_key, hashes.SHA256()))
    leaf = (builder(subject, authority_name, leaf_key)
            .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True)
            .add_extension(x509.KeyUsage(True, False, False, False, False, False, False, False, False), critical=True)
            .add_extension(x509.SubjectAlternativeName([x509.DNSName("localhost")]), critical=False)
            .add_extension(x509.ExtendedKeyUsage([ExtendedKeyUsageOID.SERVER_AUTH]), critical=False)
            .add_extension(x509.AuthorityKeyIdentifier.from_issuer_public_key(authority_key.public_key()), critical=False)
            .sign(authority_key, hashes.SHA256()))
    ca, cert, key = folder / "ca.pem", folder / "server.pem", folder / "server.key"
    private_file(ca, authority.public_bytes(serialization.Encoding.PEM))
    private_file(cert, leaf.public_bytes(serialization.Encoding.PEM))
    private_file(key, leaf_key.private_bytes(serialization.Encoding.PEM,
                 serialization.PrivateFormat.PKCS8, serialization.NoEncryption()))
    return ca, cert, key


def client_context(ca: Path | None) -> ssl.SSLContext:
    context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    context.minimum_version = ssl.TLSVersion.TLSv1_2
    if ca is not None:
        context.load_verify_locations(cafile=str(ca))
    # No system trust-store changes, automatic key logging, or verification bypass.
    return context


@contextmanager
def fixture_server(cert: Path, key: Path):
    """Serve at most one bounded request on an OS-assigned loopback port."""
    context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    context.minimum_version = ssl.TLSVersion.TLSv1_2
    context.load_cert_chain(cert, key)
    stop = threading.Event()
    observed = {"http_requests": 0, "server_error": None}
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
        listener.bind(("127.0.0.1", 0))
        listener.listen(1)
        listener.settimeout(0.1)

        def serve():
            try:
                while not stop.is_set():
                    try:
                        connection, _ = listener.accept()
                        break
                    except TimeoutError:
                        continue
                else:
                    return
                with connection:
                    connection.settimeout(TIMEOUT)
                    try:
                        secured = context.wrap_socket(connection, server_side=True)
                    except ssl.SSLError:
                        # Client rejection during the TLS handshake: no HTTP read.
                        return
                    with secured:
                        request = b""
                        while b"\r\n\r\n" not in request and len(request) < 4096:
                            chunk = secured.recv(min(1024, 4096 - len(request)))
                            if not chunk:
                                return
                            request += chunk
                        if b"\r\n\r\n" not in request:
                            return
                        observed["http_requests"] += 1
                        line = request.split(b"\r\n", 1)[0]
                        status, reason = (403, "Forbidden") if line == b"GET /restricted HTTP/1.1" else (200, "OK")
                        body = b'{"fixture":"access refused"}' if status == 403 else b'{"fixture":"service reached"}'
                        headers = f"HTTP/1.1 {status} {reason}\r\nContent-Type: application/json\r\nContent-Length: {len(body)}\r\nConnection: close\r\n\r\n"
                        secured.sendall(headers.encode("ascii") + body)
            except (OSError, ValueError):
                observed["server_error"] = "fixture_server_failed"

        thread = threading.Thread(target=serve, name="fde-loopback-fixture", daemon=True)
        thread.start()
        try:
            yield listener.getsockname()[1], observed
        finally:
            stop.set()
            thread.join(TIMEOUT + 1)
            if thread.is_alive():
                raise RuntimeError("The local fixture did not stop within its deadline.")


def exchange(case: str, port: int, ca: Path) -> dict:
    result = {"id": case, "stages": {layer: "not-run" for layer in LAYERS},
              "address": None, "expected_peer": "wrong.localhost" if case == "wrong-name" else "localhost",
              "path": "/restricted" if case == "denied" else "/ticket",
              "tls_version": None, "http_status": None, "body": None, "error": None}
    stage = "lookup"
    try:
        # The OS can answer localhost through a hosts file or cache; not a DNS capture.
        addresses = socket.getaddrinfo("localhost", port, socket.AF_INET, socket.SOCK_STREAM)
        address = addresses[0][4][0]
        if not ipaddress.ip_address(address).is_loopback or address != "127.0.0.1":
            raise ValueError("This fixture requires localhost to resolve to 127.0.0.1.")
        result["address"] = address
        result["stages"]["lookup"] = "complete"
        stage = "tcp"
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as connection:
            connection.settimeout(TIMEOUT)
            connection.connect((address, port))
            result["stages"]["tcp"] = "complete"
            stage = "tls"
            context = client_context(None if case == "untrusted" else ca)
            with context.wrap_socket(connection, server_hostname=result["expected_peer"]) as secured:
                result["stages"]["tls"] = "complete"
                result["tls_version"] = secured.version()
                stage = "http"
                request = f"GET {result['path']} HTTP/1.1\r\nHost: localhost:{port}\r\nConnection: close\r\n\r\n"
                secured.sendall(request.encode("ascii"))
                response = b""
                while len(response) < 8192:
                    part = secured.recv(min(1024, 8192 - len(response)))
                    if not part:
                        break
                    response += part
                headers, body = response.split(b"\r\n\r\n", 1)
                result["http_status"] = int(headers.split(b"\r\n", 1)[0].split()[1])
                result["body"] = json.loads(body)
                result["stages"]["http"] = "complete"
    except ssl.SSLCertVerificationError as error:
        result["error"] = {"kind": "certificate_verification", "code": error.verify_code,
                           "message": error.verify_message}
    except ConnectionRefusedError:
        result["error"] = {"kind": "connection_refused", "message": "No listener accepted this local port."}
    except TimeoutError:
        result["error"] = {"kind": "connection_timeout" if stage == "tcp" else "operation_timeout",
                           "message": f"No result arrived before the local {stage} deadline."}
    except (OSError, ValueError, IndexError):
        # Do not mislabel an unexpected local environment failure as expected success.
        result["error"] = {"kind": "environment_error", "message": f"The local attempt failed at {stage}; check the setup."}
    if result["error"]:
        result["stages"][stage] = "stopped"
    return result


def run_case(case: str) -> dict:
    if case not in CASES:
        raise ValueError("Choose one of the five fixed local cases.")
    with tempfile.TemporaryDirectory(prefix="fde-connection-") as temporary:
        ca, cert, key = certificates(Path(temporary))
        if case == "no-listener":
            # Keep this port reserved, without listen(), so no rebind race is needed. Some operating systems time out here.
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as reserved:
                reserved.bind(("127.0.0.1", 0))
                result = exchange(case, reserved.getsockname()[1], ca)
                observed = {"http_requests": 0, "server_error": None}
        else:
            with fixture_server(cert, key) as (port, observed):
                result = exchange(case, port, ca)
        result.update(observed)
    return result


def record() -> dict:
    return {
        "schema": "lockedin.first-connection/v1",
        "runtime": f"Python {platform.python_version()} · cryptography {cryptography.__version__} · {ssl.OPENSSL_VERSION}",
        "provenance": "Actual local TCP, TLS and HTTP exchanges. localhost lookup is not a captured external DNS query. Random ports and temporary certificate paths are omitted; TLS diagnostics can vary by runtime. The browser displays these recordings without opening connections.",
        "cases": [run_case(case) for case in CASES],
    }


class ConnectionTests(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.cases = {case: run_case(case) for case in CASES}

    def test_success_has_verified_tls_then_one_http_request(self):
        result = self.cases["success"]
        self.assertTrue(all(value == "complete" for value in result["stages"].values()))
        self.assertEqual((result["http_status"], result["http_requests"]), (200, 1))
        self.assertIn(result["tls_version"], ("TLSv1.2", "TLSv1.3"))
        self.assertIsNone(result["error"])

    def test_refusal_does_not_run_tls_or_http(self):
        result = self.cases["no-listener"]
        self.assertEqual(list(result["stages"].values()), ["complete", "stopped", "not-run", "not-run"])
        self.assertIn(result["error"]["kind"], ("connection_refused", "connection_timeout"))
        self.assertEqual(result["http_requests"], 0)

    def test_name_and_trust_failures_never_enter_http(self):
        for case in ("wrong-name", "untrusted"):
            with self.subTest(case=case):
                result = self.cases[case]
                self.assertEqual(list(result["stages"].values()), ["complete", "complete", "stopped", "not-run"])
                self.assertEqual(result["error"]["kind"], "certificate_verification")
                self.assertIsNone(result["http_status"])
                self.assertEqual(result["http_requests"], 0)
        self.assertNotEqual(self.cases["wrong-name"]["error"]["code"], self.cases["untrusted"]["error"]["code"])

    def test_403_is_a_response_over_a_verified_connection(self):
        result = self.cases["denied"]
        self.assertTrue(all(value == "complete" for value in result["stages"].values()))
        self.assertEqual((result["http_status"], result["http_requests"]), (403, 1))
        self.assertIsNone(result["error"])

    def test_every_connection_stays_local_and_fixture_stops(self):
        for result in self.cases.values():
            self.assertEqual(result["address"], "127.0.0.1")
            self.assertIsNone(result["server_error"])
        self.assertFalse(any(t.name == "fde-loopback-fixture" for t in threading.enumerate()))
        with self.assertRaises(ValueError):
            run_case("https://example.com")

    def test_tls_verification_cannot_be_silently_disabled(self):
        context = client_context(None)
        self.assertEqual(context.verify_mode, ssl.CERT_REQUIRED)
        self.assertTrue(context.check_hostname)
        self.assertIsNone(context.keylog_filename)
        self.assertEqual(context.cert_store_stats()["x509_ca"], 0)

    def test_temporary_keys_are_private_and_cleaned(self):
        with tempfile.TemporaryDirectory(prefix="fde-connection-test-") as temporary:
            folder = Path(temporary)
            ca, cert, key = certificates(folder)
            if os.name == "posix":
                self.assertEqual(key.stat().st_mode & 0o777, 0o600)
            with self.assertRaises(FileExistsError):
                private_file(key, b"must not replace")
            client_context(ca)
            self.assertTrue(cert.is_file())
        self.assertFalse(folder.exists())
        serialized = json.dumps(self.cases)
        self.assertNotIn("PRIVATE KEY", serialized)
        self.assertNotIn(temporary, serialized)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("command", choices=("test", "demo", "inspect"))
    parser.add_argument("--case", choices=CASES, default="success")
    args = parser.parse_args()
    if args.command == "test":
        suite = unittest.defaultTestLoader.loadTestsFromTestCase(ConnectionTests)
        raise SystemExit(0 if unittest.TextTestRunner(verbosity=2).run(suite).wasSuccessful() else 1)
    print(json.dumps(record() if args.command == "demo" else run_case(args.case), indent=2))


if __name__ == "__main__":
    main()