import hashlib, base64, json, time
from cryptography import x509
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.x509.oid import NameOID, ExtendedKeyUsageOID
import datetime, ipaddress

def generate_csr(domain: str, sans: list[str]) -> tuple[bytes, bytes]:
    """Generate EC P-256 private key + CSR. Returns (private_key_pem, csr_pem)."""
    key = ec.generate_private_key(ec.SECP256R1())

    csr = (
        x509.CertificateSigningRequestBuilder()
        .subject_name(x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, domain)]))
        .add_extension(
            x509.SubjectAlternativeName([x509.DNSName(s) for s in sans]),
            critical=False,
        )
        .sign(key, hashes.SHA256())
    )
    return (
        key.private_bytes(serialization.Encoding.PEM,
                          serialization.PrivateFormat.TraditionalOpenSSL,
                          serialization.NoEncryption()),
        csr.public_bytes(serialization.Encoding.PEM),
    )


def dns01_key_authorization(token: str, account_key_thumbprint: str) -> str:
    """Build the DNS TXT record value for dns-01 challenge."""
    key_auth = f"{token}.{account_key_thumbprint}"
    digest = hashlib.sha256(key_auth.encode()).digest()
    return base64.urlsafe_b64encode(digest).rstrip(b"=").decode()


def poll_order(acme_client, order_url: str, timeout: int = 120) -> dict:
    """Poll ACME order until valid or timeout."""
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        order = acme_client.get(order_url)
        if order["status"] == "valid":
            return order
        if order["status"] == "invalid":
            raise RuntimeError(f"ACME order invalid: {order.get('error')}")
        time.sleep(3)
    raise TimeoutError("ACME order did not complete in time")


def build_cert_bundle(leaf_pem: bytes, intermediates: list[bytes]) -> bytes:
    """Concatenate leaf + chain for nginx ssl_certificate."""
    return leaf_pem + b"\n".join(intermediates)


# ── Usage ──────────────────────────────────────────────────────────
if __name__ == "__main__":
    domain = "paddyspeaks.com"
    sans   = ["paddyspeaks.com", "www.paddyspeaks.com"]

    key_pem, csr_pem = generate_csr(domain, sans)
    print(csr_pem.decode())

    # In a real flow:
    # 1. Submit csr_pem to ACME newOrder → get challenge URLs
    # 2. Publish DNS TXT = dns01_key_authorization(token, thumbprint)
    # 3. Tell ACME server challenge is ready → poll_order(client, order_url)
    # 4. Download cert → build_cert_bundle(leaf, [intermediate_pem])
    # 5. Store key_pem in a secret (Vault, K8s Secret) — never log it

from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.x509 import ocsp
import urllib.request, ssl

def check_ocsp_status(cert_pem: bytes, issuer_pem: bytes) -> str:
    """Check OCSP status of a leaf cert against its issuer."""
    cert   = x509.load_pem_x509_certificate(cert_pem)
    issuer = x509.load_pem_x509_certificate(issuer_pem)

    # Build OCSP request
    builder = ocsp.OCSPRequestBuilder()
    builder = builder.add_certificate(cert, issuer, hashes.SHA256())
    req = builder.build()
    req_der = req.public_bytes(serialization.Encoding.DER)

    # Find OCSP URL from cert's AIA extension
    aia = cert.extensions.get_extension_for_class(x509.AuthorityInformationAccess)
    ocsp_url = next(
        ad.access_location.value
        for ad in aia.value
        if ad.access_method == x509.AuthorityInformationAccessOID.OCSP
    )

    # HTTP POST OCSP request
    http_req = urllib.request.Request(
        ocsp_url,
        data=req_der,
        headers={"Content-Type": "application/ocsp-request"},
        method="POST",
    )
    with urllib.request.urlopen(http_req, timeout=5) as resp:
        resp_der = resp.read()

    # Parse OCSP response
    ocsp_resp = ocsp.load_der_ocsp_response(resp_der)

    status_map = {
        ocsp.OCSPCertStatus.GOOD: "GOOD",
        ocsp.OCSPCertStatus.REVOKED: f"REVOKED (reason={ocsp_resp.revocation_reason})",
        ocsp.OCSPCertStatus.UNKNOWN: "UNKNOWN",
    }
    return status_map.get(ocsp_resp.certificate_status, "PARSE_ERROR")


# ── Quick cert inspection utility ─────────────────────────────────
def inspect_cert(cert_pem: bytes) -> dict:
    """Return key fields from a PEM certificate."""
    cert = x509.load_pem_x509_certificate(cert_pem)
    try:
        san = cert.extensions.get_extension_for_class(x509.SubjectAlternativeName)
        sans = san.value.get_values_for_type(x509.DNSName)
    except x509.ExtensionNotFound:
        sans = []
    return {
        "subject":      cert.subject.rfc4514_string(),
        "issuer":       cert.issuer.rfc4514_string(),
        "serial":       hex(cert.serial_number),
        "not_before":   cert.not_valid_before_utc.isoformat(),
        "not_after":    cert.not_valid_after_utc.isoformat(),
        "sans":         sans,
        "key_algo":     cert.public_key().__class__.__name__,
        "fingerprint":  cert.fingerprint(hashes.SHA256()).hex(),
    }

"""
Scan a list of hostnames, check TLS cert expiry, emit Prometheus metrics.
Deploy as a CronJob in Kubernetes or as a Prometheus exporter DaemonSet.
"""
import ssl, socket, datetime, logging
from dataclasses import dataclass
from typing import Iterator

@dataclass
class CertInfo:
    hostname: str
    port: int
    subject_cn: str
    issuer: str
    not_after: datetime.datetime
    days_remaining: int
    serial: str

def check_cert(hostname: str, port: int = 443, timeout: int = 10) -> CertInfo:
    ctx = ssl.create_default_context()
    with socket.create_connection((hostname, port), timeout=timeout) as sock:
        with ctx.wrap_socket(sock, server_hostname=hostname) as ssock:
            info = ssock.getpeercert()

    not_after = datetime.datetime.strptime(
        info["notAfter"], "%b %d %H:%M:%S %Y %Z"
    ).replace(tzinfo=datetime.timezone.utc)
    now = datetime.datetime.now(datetime.timezone.utc)
    days = (not_after - now).days

    subject = dict(x[0] for x in info["subject"])
    issuer  = dict(x[0] for x in info["issuer"])
    serial  = str(info.get("serialNumber", "unknown"))

    return CertInfo(
        hostname=hostname,
        port=port,
        subject_cn=subject.get("commonName", hostname),
        issuer=issuer.get("organizationName", "unknown"),
        not_after=not_after,
        days_remaining=days,
        serial=serial,
    )


def scan_fleet(hostnames: list[str]) -> Iterator[CertInfo]:
    for h in hostnames:
        try:
            yield check_cert(h)
        except Exception as e:
            logging.warning("cert check failed for %s: %s", h, e)


def emit_prometheus_metrics(results: list[CertInfo]) -> str:
    """Return Prometheus text format metrics."""
    lines = [
        "# HELP tls_cert_days_remaining Days until TLS certificate expires",
        "# TYPE tls_cert_days_remaining gauge",
    ]
    for r in results:
        label = f'hostname="{r.hostname}",issuer="{r.issuer}",serial="{r.serial}"'
        lines.append(f"tls_cert_days_remaining{{{label}}} {r.days_remaining}")
    return "\n".join(lines)


# ── Usage ─────────────────────────────────────────────────────────
if __name__ == "__main__":
    fleet = [
        "paddyspeaks.com", "google.com", "github.com",
        "api.stripe.com", "s3.amazonaws.com",
    ]
    results = list(scan_fleet(fleet))
    for r in sorted(results, key=lambda x: x.days_remaining):
        status = "CRITICAL" if r.days_remaining < 14 else ("WARNING" if r.days_remaining < 30 else "OK")
        print(f"[{status:8s}] {r.hostname:30s}  {r.days_remaining:3d} days  issuer={r.issuer}")
    print()
    print(emit_prometheus_metrics(results))