"""
test6_tamper_injection.py -- Test 6, the CR 3.1 leg named as open since 2026-07-31: a dedicated
integrity test, not integrity "by construction."

Why this exists
----------------
Tests 3/4/5 run over TLS, which provides record-layer integrity as a property of the channel --
but nothing in this repo had, until now, actually flipped a bit in a real record and shown it get
rejected. "By construction" (we rely on TLS) and "demonstrated" (we broke a record on purpose and
watched the AEAD tag catch it) are different tiers, and the 62443-4-2 mapping said so honestly.
This closes that gap.

What this runs
----------------
A raw TCP relay sits between a real mutual-TLS client and server (same CA/cert machinery as Tests
3-5), forwarding bytes at the TLS RECORD level -- it parses only the 5-byte record header
(content type + version + length) and never touches or needs to understand the encrypted
payload's plaintext.

  [CONTROL] Every record forwarded byte-for-byte, unmodified -> handshake completes, the app-data
            message is delivered intact. Proves the relay itself introduces no corruption of its
            own -- so anything that goes wrong in the tamper run is attributable to the tamper.
  [TAMPER]  Identical relay, except ONE BIT is flipped inside the CIPHERTEXT of the first
            Application Data record after the handshake -> the receiving TLS stack's AEAD
            authentication fails and the connection is torn down. The corrupted data is never
            delivered as if it were valid -- which is the actual empirical content of "integrity":
            not that correct data arrives, but that incorrect data provably does not arrive
            looking correct.

            Which byte, and why it matters: a TLS 1.2 AEAD record body is
            `explicit_nonce(8) || ciphertext || auth_tag(16)`, so byte 0 is the NONCE, not the
            payload. Flipping the nonce also trips the AEAD check -- but by garbling decryption,
            not by the tag catching an altered payload. CR 3.1 is about unauthorised
            MODIFICATION of transmitted information, so the flip targets a byte in the middle
            of the ciphertext. The demonstration then matches its own sentence exactly.

Run: python test6_tamper_injection.py
"""
import datetime
import socket
import ssl
import tempfile
import threading
import time
from pathlib import 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

WORK = Path(tempfile.mkdtemp(prefix="cip_test6_"))
HOST = "127.0.0.1"
PORT_SERVER = 60448
PORT_RELAY_CONTROL = 60449
PORT_RELAY_TAMPER = 60450

# TLS 1.2 AEAD (AES-GCM) record body layout, RFC 5246 §6.2.3.3 + RFC 5288 §3:
#     explicit_nonce (8) || ciphertext || auth_tag (16)
# These are named so the tamper below can target the CIPHERTEXT specifically rather than
# whatever happens to be at byte 0 (which is the nonce, not the payload).
EXPLICIT_NONCE_LEN = 8
AEAD_TAG_LEN = 16


def _wait_ready(port, timeout=5.0):
    """Poll until something is accepting on `port`. Replaces fixed sleeps, which are a
    flakiness source under load -- and a demo that intermittently fails is a demo a
    sceptical operator is right to discount."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            with socket.create_connection((HOST, port), timeout=0.25):
                return True
        except OSError:
            time.sleep(0.02)
    return False


def _key():
    return rsa.generate_private_key(public_exponent=65537, key_size=2048)


def make_ca():
    key = _key()
    name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, "TestFleet Root CA")])
    cert = (
        x509.CertificateBuilder().subject_name(name).issuer_name(name)
        .public_key(key.public_key()).serial_number(x509.random_serial_number())
        .not_valid_before(datetime.datetime.now(datetime.timezone.utc))
        .not_valid_after(datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=1))
        .add_extension(x509.BasicConstraints(ca=True, path_length=None), critical=True)
        .sign(key, hashes.SHA256())
    )
    return key, cert


def make_cert(ca_key, ca_cert, name):
    key = _key()
    subj = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, name)])
    cert = (
        x509.CertificateBuilder().subject_name(subj).issuer_name(ca_cert.subject)
        .public_key(key.public_key()).serial_number(x509.random_serial_number())
        .not_valid_before(datetime.datetime.now(datetime.timezone.utc))
        .not_valid_after(datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=1))
        .add_extension(x509.SubjectAlternativeName([x509.DNSName(name)]), critical=False)
        .sign(ca_key, hashes.SHA256())
    )
    return key, cert


def write_pem(kp, key, cp, cert):
    kp.write_bytes(key.private_bytes(
        serialization.Encoding.PEM,
        serialization.PrivateFormat.TraditionalOpenSSL,
        serialization.NoEncryption(),
    ))
    cp.write_bytes(cert.public_bytes(serialization.Encoding.PEM))


class TLSServer(threading.Thread):
    """Same role as Tests 3-5's strict endpoint: mutual TLS, then echo one message."""

    def __init__(self, port, server_cert, server_key, ca_cert):
        super().__init__(daemon=True)
        ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
        ctx.maximum_version = ssl.TLSVersion.TLSv1_2
        ctx.load_cert_chain(str(server_cert), str(server_key))
        ctx.load_verify_locations(str(ca_cert))
        ctx.verify_mode = ssl.CERT_REQUIRED
        self.ctx = ctx
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind((HOST, port))
        self.sock.listen(5)
        self.running = True

    def run(self):
        while self.running:
            try:
                raw, _ = self.sock.accept()
            except OSError:
                break
            threading.Thread(target=self._serve, args=(raw,), daemon=True).start()

    def _serve(self, raw):
        try:
            tls = self.ctx.wrap_socket(raw, server_side=True)
        except ssl.SSLError:
            try:
                raw.close()
            except OSError:
                pass
            return
        try:
            data = tls.recv(4096)
            if data:
                tls.sendall(b"ECHO:" + data)
            tls.close()
        except (OSError, ssl.SSLError):
            pass

    def stop(self):
        self.running = False
        try:
            self.sock.close()
        except OSError:
            pass


def _recv_exact(sock, n):
    buf = b""
    while len(buf) < n:
        try:
            chunk = sock.recv(n - len(buf))
        except OSError:
            return None
        if not chunk:
            return None
        buf += chunk
    return buf


def _relay_direction(src, dst, tamper, tamper_state, label):
    """Forward TLS records src->dst at the record-framing level (5-byte header: content type,
    2-byte version, 2-byte length, then that many body bytes). If `tamper`, flip one bit in the
    body of the first Application Data record (content type 0x17) seen, exactly once."""
    while True:
        header = _recv_exact(src, 5)
        if header is None:
            break
        length = int.from_bytes(header[3:5], "big")
        body = _recv_exact(src, length)
        if body is None:
            break
        content_type = header[0]
        if tamper and content_type == 0x17:
            # Target the CIPHERTEXT, not byte 0 (which is the explicit nonce). Flipping the
            # nonce also trips the AEAD check, but by garbling decryption rather than by the
            # tag catching an altered payload -- and the claim being tested is specifically
            # "unauthorised MODIFICATION of transmitted information" (62443-4-2 CR 3.1).
            # So flip a payload byte, and the demonstration matches the sentence exactly.
            lo, hi = EXPLICIT_NONCE_LEN, len(body) - AEAD_TAG_LEN
            with tamper_state["lock"]:            # both directions share this state
                mine = (not tamper_state["done"]) and hi > lo
                if mine:
                    tamper_state["done"] = True
            if mine:
                idx = (lo + hi) // 2
                body = bytearray(body)
                body[idx] ^= 0x01
                body = bytes(body)
                print(f"    [{label}] flipped one bit at ciphertext byte {idx - lo} of {hi - lo}"
                      f"  (record {length} B = nonce {EXPLICIT_NONCE_LEN}"
                      f" + ciphertext {hi - lo} + tag {AEAD_TAG_LEN})")
        try:
            dst.sendall(header + body)
        except OSError:
            break
    try:
        dst.shutdown(socket.SHUT_WR)
    except OSError:
        pass


def run_relay(listen_port, server_port, tamper):
    listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    listener.bind((HOST, listen_port))
    listener.listen(1)

    def _serve_one():
        client_sock, _ = listener.accept()
        server_sock = socket.create_connection((HOST, server_port))
        tamper_state = {"done": False, "lock": threading.Lock()}
        t1 = threading.Thread(target=_relay_direction,
                               args=(client_sock, server_sock, tamper, tamper_state, "c->s"),
                               daemon=True)
        t2 = threading.Thread(target=_relay_direction,
                               args=(server_sock, client_sock, tamper, tamper_state, "s->c"),
                               daemon=True)
        t1.start()
        t2.start()

    threading.Thread(target=_serve_one, daemon=True).start()
    return listener


def client_via_relay(relay_port, ca_cert, server_name, client_cert, client_key):
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.maximum_version = ssl.TLSVersion.TLSv1_2
    ctx.load_verify_locations(str(ca_cert))
    ctx.check_hostname = True
    ctx.verify_mode = ssl.CERT_REQUIRED
    ctx.load_cert_chain(str(client_cert), str(client_key))
    try:
        with socket.create_connection((HOST, relay_port), timeout=3) as raw:
            with ctx.wrap_socket(raw, server_hostname=server_name) as tls:
                tls.sendall(b"hello, device-a")
                data = tls.recv(1024)
                return data.decode(errors="replace") if data else "REJECTED (connection closed, no data)"
    except ssl.SSLError as e:
        return f"REJECTED (TLS integrity/handshake error): {e.__class__.__name__}: {e}"
    except OSError as e:
        return f"REJECTED (connection error): {e}"


def main():
    print("[*] CA + server cert (device-a) + client cert (device-a-client).\n")
    ca_key, ca_cert = make_ca()
    s_key, s_cert = make_cert(ca_key, ca_cert, "device-a")
    c_key, c_cert = make_cert(ca_key, ca_cert, "device-a-client")
    ca_p = WORK / "ca.pem"
    ca_p.write_bytes(ca_cert.public_bytes(serialization.Encoding.PEM))
    s_kp, s_cp = WORK / "s.key", WORK / "s.pem"
    write_pem(s_kp, s_key, s_cp, s_cert)
    c_kp, c_cp = WORK / "c.key", WORK / "c.pem"
    write_pem(c_kp, c_key, c_cp, c_cert)

    server = TLSServer(PORT_SERVER, s_cp, s_kp, ca_p)
    server.start()
    if not _wait_ready(PORT_SERVER):
        print("[!] server never came up on", PORT_SERVER)
        return

    try:
        # No sleep needed after run_relay(): it binds and listens synchronously before
        # returning, so a client connecting immediately is held in the listen backlog.
        print("[CONTROL] relay forwards every record byte-for-byte, unmodified:")
        run_relay(PORT_RELAY_CONTROL, PORT_SERVER, tamper=False)
        r_control = client_via_relay(PORT_RELAY_CONTROL, ca_p, "device-a", c_cp, c_kp)
        print("  ->", r_control)

        print("\n[TAMPER] identical relay, except one bit flips in the ciphertext of the first Application Data record:")
        run_relay(PORT_RELAY_TAMPER, PORT_SERVER, tamper=True)
        r_tamper = client_via_relay(PORT_RELAY_TAMPER, ca_p, "device-a", c_cp, c_kp)
        print("  ->", r_tamper)

        print("\n" + "=" * 72)
        control_ok = "ECHO:hello, device-a" in r_control
        tamper_ok = "REJECTED" in r_tamper
        if control_ok and tamper_ok:
            print("[!] The control proves the relay itself is transparent -- the message arrives")
            print("    intact when nothing is altered. The tamper run flips a single bit inside the")
            print("    CIPHERTEXT of an encrypted record whose plaintext the relay never sees (not")
            print("    the nonce, not the tag -- the payload), and the receiving TLS stack's AEAD")
            print("    authentication check catches it and tears down the connection -- the")
            print("    corrupted data is never delivered as if it were valid.")
            print("    That is CR 3.1's integrity clause, demonstrated rather than assumed from")
            print("    'we used TLS.'")
        else:
            print(f"[?] Unexpected: control_ok={control_ok} tamper_ok={tamper_ok}")
            print(f"    r_control={r_control!r}")
            print(f"    r_tamper={r_tamper!r}")
    finally:
        server.stop()


if __name__ == "__main__":
    main()