#!/usr/bin/env python
"""One TLS handshake per domain: who issued the certificate, and how long does it live?

Method and verdict rules are in README.md, written before this ran.

    python probe.py --domains ../llms-txt/domains.csv --out certs.csv

Two handshakes per domain: the first with verification off (the certificate
that is *served* is the subject, including an expired or mismatched one), the
second with the system trust store on, only to record whether it would verify.
"""

import argparse
import concurrent.futures as futures
import csv
import datetime as dt
import socket
import ssl
import sys

from cryptography import x509
from cryptography.hazmat.primitives.asymmetric import ec, ed25519, rsa

TIMEOUT = 10


def _name_part(name, oid):
    try:
        vals = name.get_attributes_for_oid(oid)
    except Exception:
        return ""
    return vals[0].value if vals else ""


def key_of(cert):
    pub = cert.public_key()
    if isinstance(pub, rsa.RSAPublicKey):
        return "RSA", pub.key_size
    if isinstance(pub, ec.EllipticCurvePublicKey):
        return "ECDSA", pub.curve.key_size
    if isinstance(pub, ed25519.Ed25519PublicKey):
        return "Ed25519", 256
    return type(pub).__name__, 0


def verifies(domain):
    """A second handshake with the system trust store. Returns (bool, reason)."""
    ctx = ssl.create_default_context()
    try:
        with socket.create_connection((domain, 443), TIMEOUT) as raw:
            with ctx.wrap_socket(raw, server_hostname=domain):
                return True, ""
    except ssl.SSLCertVerificationError as exc:
        return False, f"verify:{exc.verify_message or exc.reason}"
    except Exception as exc:                       # noqa: BLE001 — recorded, not swallowed
        return False, f"{type(exc).__name__}:{exc}"


def probe(rank, domain):
    row = {"rank": rank, "domain": domain, "ok": 0, "error": "",
           "issuer_org": "", "issuer_cn": "", "subject_cn": "",
           "not_before": "", "not_after": "", "lifetime_days": "",
           "key_alg": "", "key_bits": "", "tls": "", "cipher": "",
           "verified": "", "verify_error": ""}
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    try:
        with socket.create_connection((domain, 443), TIMEOUT) as raw:
            with ctx.wrap_socket(raw, server_hostname=domain) as sock:
                der = sock.getpeercert(binary_form=True)
                row["tls"] = sock.version() or ""
                cipher = sock.cipher()
                row["cipher"] = cipher[0] if cipher else ""
    except Exception as exc:                       # noqa: BLE001
        row["error"] = f"{type(exc).__name__}:{exc}"[:160]
        return row
    if not der:
        row["error"] = "no_certificate"
        return row
    try:
        cert = x509.load_der_x509_certificate(der)
    except Exception as exc:                       # noqa: BLE001
        row["error"] = f"parse:{type(exc).__name__}"
        return row
    row["issuer_org"] = _name_part(cert.issuer, x509.oid.NameOID.ORGANIZATION_NAME)
    row["issuer_cn"] = _name_part(cert.issuer, x509.oid.NameOID.COMMON_NAME)
    row["subject_cn"] = _name_part(cert.subject, x509.oid.NameOID.COMMON_NAME)
    nb, na = cert.not_valid_before_utc, cert.not_valid_after_utc
    row["not_before"] = nb.date().isoformat()
    row["not_after"] = na.date().isoformat()
    row["lifetime_days"] = (na - nb).days
    alg, bits = key_of(cert)
    row["key_alg"], row["key_bits"] = alg, bits
    ok, why = verifies(domain)
    row["verified"] = 1 if ok else 0
    row["verify_error"] = why[:160]
    row["ok"] = 1
    return row


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--domains", required=True, help="csv with rank,domain columns")
    ap.add_argument("--out", required=True)
    ap.add_argument("--limit", type=int, default=1000)
    ap.add_argument("--workers", type=int, default=32)
    args = ap.parse_args(argv)

    want = []
    with open(args.domains, newline="") as fh:
        for r in csv.DictReader(fh):
            want.append((int(r["rank"]), r["domain"]))
            if len(want) >= args.limit:
                break

    rows = []
    with futures.ThreadPoolExecutor(args.workers) as pool:
        jobs = {pool.submit(probe, rank, d): d for rank, d in want}
        for n, fut in enumerate(futures.as_completed(jobs), 1):
            rows.append(fut.result())
            if n % 100 == 0:
                print(f"  {n}/{len(want)}", file=sys.stderr, flush=True)
    rows.sort(key=lambda r: r["rank"])
    with open(args.out, "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)
    ok = sum(r["ok"] for r in rows)
    print(f"{ok}/{len(rows)} handshakes, written to {args.out}")
    print("as of", dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())