"""
test5_rotation.py -- Test 5, the CR 1.8 lifecycle leg Test 3/4 did not close: ROTATION.

Why this exists
----------------
Test 3 proved per-device IDENTITY (uniqueness). Test 4 proved REVOCATION (a credential can be
taken back). Neither proves ROTATION -- issuing a REPLACEMENT credential for the same identity
and RETIRING the prior one, which is what CR 1.8 actually asks for ("issue a replacement
credential and retire the prior one"). Named as the one open item in the README / 62443-4-2
mapping since 2026-07-31. This closes it.

The control this needs (same discipline as Test 3's case 4): does merely ISSUING a new credential
automatically retire the old one? If rotation "just worked" from issuance alone, an attacker who'd
stolen the old (still-valid) credential would keep access after a legitimate rotation -- silently.
The control below shows re-issuance and retirement are two SEPARATE actions; skipping the second
leaves the old credential live.

What this runs
----------------
  [1] v1 credential for 'engineer-1'                     -> GRANTED  (before rotation)
  [2] v2 credential issued for the SAME identity          -> GRANTED  (re-issuance works)
  [3] v1 STILL presented, after v2 exists, BEFORE v1 is explicitly retired
                                                           -> GRANTED  <-- THE CONTROL
      (proves issuing v2 did NOT automatically retire v1 -- they are independent facts)
  [4] v1 explicitly retired (added to the CRL)            -> DENIED
  [5] v2 presented again, unaffected by v1's retirement   -> GRANTED  (continuity: the identity
      never lost access during the transition)

Run: python test5_rotation.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_test5_"))
HOST, PORT = "127.0.0.1", 60447


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):
    """A fresh keypair + cert for `name`. Calling this twice for the same name models
    re-issuance: same identity, new key material, new serial."""
    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 build_crl(ca_key, ca_cert, revoked_serials):
    b = (
        x509.CertificateRevocationListBuilder().issuer_name(ca_cert.subject)
        .last_update(datetime.datetime.now(datetime.timezone.utc))
        .next_update(datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=1))
    )
    for s in revoked_serials:
        rc = (
            x509.RevokedCertificateBuilder().serial_number(s)
            .revocation_date(datetime.datetime.now(datetime.timezone.utc)).build()
        )
        b = b.add_revoked_certificate(rc)
    return b.sign(ca_key, hashes.SHA256())


def revoked_now(crl_path):
    crl = x509.load_pem_x509_crl(Path(crl_path).read_bytes())
    return {rc.serial_number for rc in crl}


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 Endpoint(threading.Thread):
    def __init__(self, server_cert, server_key, ca_cert, crl_path):
        super().__init__(daemon=True)
        self.crl_path = crl_path
        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
            try:
                tls = self.ctx.wrap_socket(raw, server_side=True)
            except ssl.SSLError:
                try:
                    raw.close()
                except OSError:
                    pass
                continue
            try:
                der = tls.getpeercert(binary_form=True)
                serial = x509.load_der_x509_certificate(der).serial_number
                if serial in revoked_now(self.crl_path):
                    tls.sendall(b"ACCESS DENIED (credential REVOKED -- serial is on the current CRL)")
                else:
                    tls.sendall(b"ACCESS GRANTED (CA-valid, unexpired, and not revoked)")
                tls.close()
            except (OSError, ssl.SSLError):
                pass

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


def client(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, PORT), timeout=3) as raw:
            with ctx.wrap_socket(raw, server_hostname=server_name) as tls:
                return tls.recv(1024).decode()
    except ssl.SSLError as e:
        return f"REJECTED (handshake): {e}"


def main():
    print("[*] CA + endpoint cert (device-a) + FIRST credential (engineer-1, v1).\n")
    ca_key, ca_cert = make_ca()
    s_key, s_cert = make_cert(ca_key, ca_cert, "device-a")
    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)

    v1_key, v1_cert = make_cert(ca_key, ca_cert, "engineer-1")
    v1_kp, v1_cp = WORK / "v1.key", WORK / "v1.pem"
    write_pem(v1_kp, v1_key, v1_cp, v1_cert)

    crl_p = WORK / "crl.pem"
    crl_p.write_bytes(build_crl(ca_key, ca_cert, []).public_bytes(serialization.Encoding.PEM))

    ep = Endpoint(s_cp, s_kp, ca_p, crl_p)
    ep.start()
    time.sleep(0.5)
    try:
        r1 = client(ca_p, "device-a", v1_cp, v1_kp)
        print("  [1] v1 credential, before rotation            ->", r1)

        print("\n[*] Rotating: issuing a NEW (v2) credential for the SAME identity 'engineer-1'.")
        v2_key, v2_cert = make_cert(ca_key, ca_cert, "engineer-1")
        v2_kp, v2_cp = WORK / "v2.key", WORK / "v2.pem"
        write_pem(v2_kp, v2_key, v2_cp, v2_cert)

        r2 = client(ca_p, "device-a", v2_cp, v2_kp)
        print("  [2] v2 credential, freshly issued              ->", r2)

        r3 = client(ca_p, "device-a", v1_cp, v1_kp)
        print("  [3] v1 STILL presented, NOT yet retired        ->", r3, "  <-- THE CONTROL")

        print("\n[*] Explicitly retiring v1: adding its serial to the CRL.")
        crl_p.write_bytes(
            build_crl(ca_key, ca_cert, [v1_cert.serial_number]).public_bytes(serialization.Encoding.PEM)
        )

        r4 = client(ca_p, "device-a", v1_cp, v1_kp)
        print("  [4] v1 presented, AFTER explicit retirement    ->", r4)

        r5 = client(ca_p, "device-a", v2_cp, v2_kp)
        print("  [5] v2 presented again, unaffected             ->", r5)

        print("\n" + "=" * 72)
        ok = ("GRANTED" in r1 and "GRANTED" in r2 and "GRANTED" in r3
              and "DENIED" in r4 and "REVOKED" in r4 and "GRANTED" in r5)
        if ok:
            print("[!] Case 3 is the control: issuing v2 did NOT retire v1 -- both were valid")
            print("    simultaneously until v1 was EXPLICITLY revoked (case 4). Rotation is two")
            print("    actions, not one: re-issue (case 2) AND retire (case 4). Case 5 shows the")
            print("    identity never lost access during the transition -- that's what makes this")
            print("    a rotation, not just a swap. CR 1.8's 'issue a replacement and retire the")
            print("    prior one' is now demonstrated, not merely enabled by the same PKI.")
        else:
            print(f"[?] Unexpected: r1={r1!r} r2={r2!r} r3={r3!r} r4={r4!r} r5={r5!r}")
    finally:
        ep.stop()


if __name__ == "__main__":
    main()