#!/usr/bin/env python3
"""
verify-cert-compression.py — reproduce every number on
https://pqcrypta.com/cert-compression/ against any TLS 1.3 server.

Nothing here is privileged: it pulls the server's certificate chain, rebuilds
the TLS 1.3 Certificate message exactly as RFC 8446 4.4.2 specifies, and
compresses it the way RFC 8879 does. If the server compresses its own chain,
the byte count this prints is the byte count on the wire.

    python3 verify-cert-compression.py pqcrypta.com
    python3 verify-cert-compression.py example.org --json

Requires openssl in PATH. Brotli and zstandard are optional; install them to
see algorithms 2 and 3 alongside zlib.
"""

import argparse
import json
import re
import subprocess
import sys
import zlib

# RFC 8879 registry: the algorithm identifiers carried in compress_certificate.
ALGORITHMS = {1: "zlib", 2: "brotli", 3: "zstd"}


def run_openssl(args, timeout, stdin=b""):
    """Run openssl, turning the three predictable failures into clean exits."""
    try:
        return subprocess.run(["openssl"] + args, input=stdin,
                              capture_output=True, timeout=timeout)
    except FileNotFoundError:
        sys.exit("openssl not found in PATH. This script needs the openssl "
                 "command-line tool; install it and try again.")
    except subprocess.TimeoutExpired:
        sys.exit(f"openssl timed out after {timeout}s. The host may be "
                 f"unreachable, firewalled, or not speaking TLS on that port. "
                 f"Raise the limit with --timeout if the server is just slow.")


def fetch_chain(host, port, timeout):
    """Return the server's certificate chain as a list of DER byte strings."""
    proc = run_openssl(["s_client", "-connect", f"{host}:{port}",
                        "-servername", host, "-tls1_3", "-showcerts"], timeout)
    pems = re.findall(
        rb"-----BEGIN CERTIFICATE-----.*?-----END CERTIFICATE-----",
        proc.stdout, re.S,
    )
    if not pems:
        err = proc.stderr.decode(errors="replace")
        if "Name or service not known" in err or "nodename nor servname" in err:
            hint = "that hostname does not resolve"
        elif "Connection refused" in err:
            hint = "nothing is listening on that port"
        elif "no protocols available" in err or "unsupported protocol" in err:
            hint = "the server does not speak TLS 1.3, which RFC 8879 requires"
        else:
            hint = "no TLS 1.3 handshake completed"
        sys.exit(f"could not read a certificate chain from {host}:{port}: "
                 f"{hint}.\n\nopenssl said:\n"
                 f"{err.strip()[:400] or '(nothing)'}")
    ders = []
    for pem in pems:
        der = run_openssl(["x509", "-inform", "PEM", "-outform", "DER"],
                          timeout, stdin=pem).stdout
        ders.append(der)
    return ders


def observed_compression(host, port, timeout):
    """What the server actually did, read off the handshake.

    Returns (algorithm, uncompressed, compressed) or None. The caller must
    also check `offered` — a None here only means something if we asked.
    """
    proc = run_openssl(["s_client", "-connect", f"{host}:{port}",
                        "-servername", host, "-tls1_3", "-trace"], timeout)
    text = proc.stdout.decode(errors="replace") + proc.stderr.decode(errors="replace")

    # Positive control: did *we* advertise compress_certificate at all? Without
    # this, "no compression" is indistinguishable from "never asked", and a
    # broken probe reads as a clean negative on every host it touches.
    offered = re.search(r"compress_certificate\(27\).*?\n\s+(\w+)", text, re.S)

    m = re.search(
        r"CompressedCertificate, Length=\d+\s*\n"
        r"\s*Compression type=(\w+).*?\n"
        r"\s*Uncompressed length=(\d+).*?\n"
        r"\s*Compressed length=(\d+)",
        text,
    )
    result = None
    if m:
        result = (m.group(1), int(m.group(2)), int(m.group(3)))
    return offered.group(1) if offered else None, result


def certificate_message(ders):
    """Rebuild the TLS 1.3 Certificate message body (RFC 8446 4.4.2).

        opaque certificate_request_context<0..2^8-1>    -- 1 length byte
        CertificateEntry certificate_list<0..2^24-1>    -- 3 length bytes
      each entry:
        opaque cert_data<1..2^24-1>                     -- 3 length bytes
        Extension extensions<0..2^16-1>                 -- 2 length bytes
    """
    def u24(n):
        return bytes([(n >> 16) & 0xFF, (n >> 8) & 0xFF, n & 0xFF])

    entries = b"".join(u24(len(d)) + d + b"\x00\x00" for d in ders)
    return b"\x00" + u24(len(entries)) + entries


def zlib_compress(data):
    """zlib as RFC 8879 requires: the zlib wrapper (RFC 1950), not raw DEFLATE."""
    c = zlib.compressobj(6, zlib.DEFLATED, 15)
    return c.compress(data) + c.flush()


def high_entropy_spans(pems_dir_certs, message):
    """Locate signature and public-key bytes inside the message.

    These are cryptographically random and cannot compress; finding them is
    what explains the gap between the headline ratio and what the compressor
    achieves on the material it can actually work with.
    """
    spans = []
    for der in pems_dir_certs:
        text = run_openssl(["x509", "-inform", "DER", "-noout", "-text"],
                           30, stdin=der).stdout.decode(errors="replace")
        for pattern in (r"Signature Value:\n((?:\s+[0-9a-f:]+\n)+)",
                        r"pub:\n((?:\s+[0-9a-f:]+\n)+)"):
            m = re.search(pattern, text)
            if not m:
                continue
            blob = bytes.fromhex(re.sub(r"[^0-9a-f]", "", m.group(1)))
            pos = message.find(blob)
            if pos >= 0:
                spans.append((pos, pos + len(blob)))
    return sorted(spans)


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("host")
    ap.add_argument("-p", "--port", type=int, default=443)
    ap.add_argument("-t", "--timeout", type=int, default=25)
    ap.add_argument("--json", action="store_true", help="machine-readable output")
    args = ap.parse_args()

    ders = fetch_chain(args.host, args.port, args.timeout)
    message = certificate_message(ders)
    compressed = zlib_compress(message)
    offered, observed = observed_compression(args.host, args.port, args.timeout)

    # Split the message into material that can compress and material that cannot.
    spans = high_entropy_spans(ders, message)
    entropy_bytes = sum(b - a for a, b in spans)
    structure, prev = b"", 0
    for a, b in spans:
        structure += message[prev:a]
        prev = b
    structure += message[prev:]
    structure_z = zlib_compress(structure)
    entropy_blob = b"".join(message[a:b] for a, b in spans)
    entropy_z = zlib_compress(entropy_blob) if entropy_blob else b""

    algos = {"zlib": len(compressed)}
    try:
        import brotli
        algos["brotli"] = len(brotli.compress(message, quality=11))
    except ImportError:
        pass
    try:
        import zstandard
        algos["zstd"] = len(zstandard.ZstdCompressor(level=19).compress(message))
    except ImportError:
        pass

    report = {
        "host": args.host,
        "chain": [len(d) for d in ders],
        "der_total": sum(len(d) for d in ders),
        "framing": len(message) - sum(len(d) for d in ders),
        "message": len(message),
        "algorithms": algos,
        "saved_payload": len(message) - len(compressed),
        "saved_wire": len(message) - (len(compressed) + 8),
        "incompressible": {
            "bytes": entropy_bytes,
            "share": round(100 * entropy_bytes / len(message), 1),
            "zlib": len(entropy_z),
        },
        "structure": {"bytes": len(structure), "zlib": len(structure_z)},
        "offered_by_probe": offered,
        "observed": None,
    }
    if observed:
        report["observed"] = {"algorithm": observed[0],
                              "uncompressed": observed[1],
                              "compressed": observed[2]}

    if args.json:
        print(json.dumps(report, indent=2))
        return

    print(f"\n  {args.host}:{args.port}\n")
    print(f"  chain of {len(ders)} certificates")
    for i, d in enumerate(ders, 1):
        subj = run_openssl(["x509", "-inform", "DER", "-noout", "-subject"],
                           30, stdin=d).stdout.decode(errors="replace").strip()
        subj = subj.replace("subject=", "").strip()
        print(f"    {i}. {len(d):5d} B  {subj[:58]}")
    print(f"\n    DER total          {report['der_total']:6d} B")
    print(f"    TLS framing        {report['framing']:6d} B")
    print(f"    Certificate msg    {report['message']:6d} B")

    print(f"\n  compressed (this script, offline)")
    for name, size in algos.items():
        ident = next(k for k, v in ALGORITHMS.items() if v == name)
        print(f"    {name:7s} ({ident})      {size:6d} B   "
              f"ratio {len(message)/size:.3f}:1   saved {len(message)-size}")

    print(f"\n  where the limit comes from")
    print(f"    signatures + keys  {entropy_bytes:6d} B  "
          f"({report['incompressible']['share']}% of the message)")
    if entropy_blob:
        print(f"      compressed       {len(entropy_z):6d} B   "
              f"ratio {len(entropy_blob)/len(entropy_z):.3f}:1  "
              f"{'-- it GROWS' if len(entropy_z) >= len(entropy_blob) else ''}")
    print(f"    everything else    {len(structure):6d} B")
    print(f"      compressed       {len(structure_z):6d} B   "
          f"ratio {len(structure)/len(structure_z):.3f}:1")

    print(f"\n  on the wire")
    if offered:
        print(f"    probe offered      {offered}")
    else:
        print(f"    probe offered      NOTHING -- result below proves nothing")
    if observed:
        alg, unc, comp = observed
        print(f"    server sent        CompressedCertificate, {alg}")
        print(f"    {unc} -> {comp} B  (ratio {unc/comp:.6f}:1)")
        if comp == len(compressed):
            print(f"    reproduced exactly: this script got {len(compressed)} B too")
        else:
            print(f"    this script got {len(compressed)} B "
                  f"(differs by {abs(comp-len(compressed))} -- "
                  f"compressor settings differ)")
    elif offered:
        print(f"    server sent        a plain Certificate message")
        print(f"    it was offered {offered} and declined: a finding about its")
        print(f"    configuration, not a gap in this measurement.")
    print()


if __name__ == "__main__":
    main()
