"""
test4_revocation.py -- the lifecycle leg Test 3 did not demonstrate: REVOCATION.

Test 3 proved per-device IDENTITY binding (uniqueness). Uniqueness is NOT
revocability: a unique credential you cannot take back is a key you can never
recall once a device is compromised or an engineer leaves. IEC 62443-4-2 CR 1.8
(PKI) and CR 1.9 (public-key auth) both require revocation-status checking; Test 3
did not show it. This does.

A real CRL (Certificate Revocation List) is signed by the CA. The endpoint checks
each presented credential's serial against the *current* CRL and denies a revoked
one -- even though that credential is still genuinely CA-signed and unexpired.
That "still valid, but revoked -> denied" case is the empirical content of the
revocation clause. (Production would use OCSP or CRL distribution points at the TLS
layer; this models the check explicitly so the logic is legible, exactly as Test 3
modeled the identity check explicitly.)

  [1] engineer-1's credential, not revoked   -> GRANTED
  [2] the SAME credential after revocation    -> DENIED (revoked)   <-- THE CONTROL

[2] is what makes revocation real: the credential did not expire, was not tampered,
still chains to the CA -- and is refused, because the CRL now lists its serial.
"""
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_test4_"))
HOST, PORT = "127.0.0.1", 60446


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 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):   # re-read: always the CURRENT CRL
                    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) + one client credential (engineer-1).\n")
    ca_key, ca_cert = make_ca()
    s_key, s_cert = make_cert(ca_key, ca_cert, "device-a")      # endpoint's TLS identity
    c_key, c_cert = make_cert(ca_key, ca_cert, "engineer-1")    # the credential we will revoke
    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)
    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", c_cp, c_kp)
        print("  [1] engineer-1, NOT revoked            ->", r1)

        crl_p.write_bytes(build_crl(ca_key, ca_cert, [c_cert.serial_number]).public_bytes(serialization.Encoding.PEM))
        print("  [*] engineer-1's serial added to the CRL (revoked) -- same cert, unchanged, unexpired.")

        r2 = client(ca_p, "device-a", c_cp, c_kp)
        print("  [2] engineer-1, AFTER revocation       ->", r2, "  <-- THE CONTROL")

        print("\n" + "=" * 72)
        if "GRANTED" in r1 and "DENIED" in r2 and "REVOKED" in r2:
            print("[!] The identical, still-valid, unexpired, CA-signed credential is")
            print("    GRANTED before revocation and DENIED after -- solely because the")
            print("    CRL now lists its serial. That is revocability: the ability to take")
            print("    a credential BACK. Test 3 proved the key is unique; this proves it")
            print("    can be recalled -- the CR 1.8 / 1.9 revocation clause, demonstrated.")
        else:
            print(f"[?] Unexpected: r1={r1!r} r2={r2!r}")
    finally:
        ep.stop()


if __name__ == "__main__":
    main()