#!/usr/bin/python3
"""Verify Squid's ssl-bump certificate mimicking against the installed OpenSSL.

This exercises src/ssl/gadgets.cc and
src/security/cert_generators/file/certificate_db.cc through the
security_file_certgen helper, which is the only part of squid that
manipulates ASN1_STRING internals and X509 name accessors directly.

It serves as a regression test for the OpenSSL 4 port,
where struct asn1_string_st became opaque and the X509 getters became const.
Two of those call sites transfer ownership of an ASN1_item_i2d() buffer into an
ASN1_OCTET_STRING; getting them wrong silently drops the extension rather than
failing to build, so a compile-only check is not enough.

Checks, per generated certificate:
  * authorityKeyIdentifier is present and equals the signing CA's
    subjectKeyIdentifier            -> mimicAuthorityKeyId(), ASN1_STRING_set0()
  * keyUsage / extendedKeyUsage / basicConstraints are mimicked
                                    -> mimicExtensions(), X509_get_ext() const
  * keyEncipherment is force-added when mimicking an EC certificate, and is
    NOT added when mimicking an RSA one
                                    -> mimicExtensions() keyusage fixup,
                                       X509_EXTENSION_set_data()
  * the subject CN is the requested one and appears as a DNS SAN
                                    -> replaceCommonName(),
                                       addAltNameWithSubjectCn()
  * the on-disk database records a usable expiry date and subject, and a
    repeated request is served from the database
                                    -> certificate_db.cc ASN1_STRING_get0_data(),
                                       OneLineSummary(), InRamCertificateDbKey()
"""

import os
import re
import shutil
import subprocess
import sys
import tempfile

CERTGEN = "/usr/lib/squid/security_file_certgen"
CA_CN = "Squid autopkgtest CA"

failures = []
checks = 0


def check(ok, what, detail=""):
    global checks
    checks += 1
    if ok:
        print("PASS: %s" % what)
    else:
        print("FAIL: %s%s" % (what, ("\n      " + detail) if detail else ""))
        failures.append(what)


def run(*argv, **kwargs):
    """Run a command, fail the whole test loudly if it errors."""
    return subprocess.run(
        argv, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
        universal_newlines=True, **kwargs).stdout


# ---------------------------------------------------------------- fixtures ---

CA_CNF = """
[req]
distinguished_name = dn
prompt = no
[dn]
CN = {cn}
[v3_ca]
basicConstraints = critical,CA:TRUE
keyUsage = critical,keyCertSign,cRLSign
subjectKeyIdentifier = hash
""".format(cn=CA_CN)

# Deliberately no keyEncipherment: squid must add it back when it replaces the
# origin's EC key with an RSA one, because NSS-based clients require it.
ORIGIN_CNF = """
[req]
distinguished_name = dn
prompt = no
[dn]
CN = {cn}
[v3_origin]
basicConstraints = critical,CA:FALSE
keyUsage = critical,digitalSignature,keyAgreement
extendedKeyUsage = serverAuth,clientAuth
subjectKeyIdentifier = hash
authorityKeyIdentifier = keyid,issuer
subjectAltName = DNS:{cn},IP:192.0.2.10
"""


def make_ca(workdir):
    cnf = os.path.join(workdir, "ca.cnf")
    key = os.path.join(workdir, "ca.key")
    crt = os.path.join(workdir, "ca.pem")
    with open(cnf, "w") as f:
        f.write(CA_CNF)
    run("openssl", "req", "-x509", "-config", cnf, "-extensions", "v3_ca",
        "-newkey", "rsa:2048", "-nodes", "-keyout", key, "-out", crt,
        "-days", "3650", "-sha256")
    return crt, key


def make_origin(workdir, cn, keytype):
    """Create a cert to be mimicked. keytype is 'ec' or 'rsa'."""
    cnf = os.path.join(workdir, "%s.cnf" % keytype)
    key = os.path.join(workdir, "%s.key" % keytype)
    csr = os.path.join(workdir, "%s.csr" % keytype)
    crt = os.path.join(workdir, "%s.pem" % keytype)
    with open(cnf, "w") as f:
        f.write(ORIGIN_CNF.format(cn=cn))
    if keytype == "ec":
        run("openssl", "ecparam", "-name", "prime256v1", "-genkey", "-noout",
            "-out", key)
    else:
        run("openssl", "genrsa", "-out", key, "2048")
    run("openssl", "req", "-new", "-config", cnf, "-key", key, "-out", csr)
    run("openssl", "x509", "-req", "-in", csr, "-CA", CA_CRT, "-CAkey", CA_KEY,
        "-CAcreateserial", "-out", crt, "-days", "365", "-sha256",
        "-extfile", cnf, "-extensions", "v3_origin")
    return crt


# ------------------------------------------------------- certgen protocol ---

class CertGen(object):
    """Speaks the security_file_certgen helper protocol.

    Request:  "new_certificate <body length> <body>"
    Body:     "host=<cn>\\nSign=...\\nSignHash=...\\n" followed by the signing
              CA certificate, the CA private key, and finally the certificate
              to mimic, all PEM.
    Reply:    "OK <body length> <cert PEM><key PEM>" terminated by \\x01.
    """

    def __init__(self, db_path=None):
        argv = [CERTGEN]
        if db_path:
            argv += ["-s", db_path, "-M", "4MB"]
        self.proc = subprocess.Popen(
            argv, stdin=subprocess.PIPE, stdout=subprocess.PIPE)

    def request(self, host, mimic_pem):
        body = "host=%s\nSign=signTrusted\nSignHash=sha256\n" % host
        body += read(CA_CRT) + read(CA_KEY) + read(mimic_pem)
        raw = body.encode()
        msg = b"new_certificate " + str(len(raw)).encode() + b" " + raw
        self.proc.stdin.write(msg)
        self.proc.stdin.flush()

        reply = b""
        while not reply.endswith(b"\x01"):
            chunk = self.proc.stdout.read(1)
            if not chunk:
                raise RuntimeError(
                    "security_file_certgen closed its output early; "
                    "partial reply: %r" % reply[:400])
            reply += chunk
        reply = reply[:-1].decode()

        code, _, rest = reply.partition(" ")
        if code != "OK":
            raise RuntimeError("certgen refused the request: %s" % reply[:400])
        return rest.partition(" ")[2]

    def close(self):
        self.proc.stdin.close()
        self.proc.wait()


def read(path):
    with open(path) as f:
        return f.read()


# ------------------------------------------------------------ x509 helpers ---

def ext(pem_text, name):
    """Return the printed body of one X509v3 extension, or None."""
    with tempfile.NamedTemporaryFile("w", suffix=".pem", delete=False) as f:
        f.write(pem_text)
        path = f.name
    try:
        out = subprocess.run(
            ["openssl", "x509", "-in", path, "-noout", "-ext", name],
            stdout=subprocess.PIPE, stderr=subprocess.DEVNULL,
            universal_newlines=True).stdout
    finally:
        os.unlink(path)
    # Drop the "X509v3 Foo: [critical]" header line, keep the value.
    lines = [ln.strip() for ln in out.splitlines() if ln.strip()]
    if not lines or ":" not in lines[0]:
        return None
    body = " ".join(lines[1:]) if len(lines) > 1 else ""
    return body or None


def keyid(pem_text, name):
    """Extract the colon-separated hex key identifier from an extension."""
    body = ext(pem_text, name)
    if not body:
        return None
    m = re.search(r"(?:[0-9A-Fa-f]{2}:){3,}[0-9A-Fa-f]{2}", body)
    return m.group(0).upper() if m else None


def text(pem_text, *args):
    with tempfile.NamedTemporaryFile("w", suffix=".pem", delete=False) as f:
        f.write(pem_text)
        path = f.name
    try:
        return run("openssl", "x509", "-in", path, "-noout", *args)
    finally:
        os.unlink(path)


# ------------------------------------------------------------------ tests ---

def check_generated(pem, host, expect_key_encipherment, label):
    ca_skid = keyid(read(CA_CRT), "subjectKeyIdentifier")
    check(ca_skid is not None, "%s: test CA has a subjectKeyIdentifier" % label)

    akid = keyid(pem, "authorityKeyIdentifier")
    check(akid is not None,
          "%s: generated cert has an authorityKeyIdentifier" % label,
          "mimicAuthorityKeyId() produced nothing; the ASN1_OCTET_STRING "
          "holding the DER almost certainly came out empty")
    check(akid == ca_skid,
          "%s: authorityKeyIdentifier matches the CA subjectKeyIdentifier"
          % label, "got %s, want %s" % (akid, ca_skid))

    ku = ext(pem, "keyUsage")
    check(ku is not None, "%s: keyUsage was mimicked" % label)
    check(ext(pem, "extendedKeyUsage") is not None,
          "%s: extendedKeyUsage was mimicked" % label)
    check(ext(pem, "basicConstraints") is not None,
          "%s: basicConstraints was mimicked" % label)

    has_ke = bool(ku and "Key Encipherment" in ku)
    if expect_key_encipherment:
        check(has_ke,
              "%s: keyEncipherment added to the mimicked EC keyUsage" % label,
              "keyUsage is %r; X509_EXTENSION_set_data() did not take effect, "
              "so NSS/Firefox clients would reject this certificate" % ku)
    else:
        check(not has_ke,
              "%s: keyEncipherment not added when mimicking an RSA cert"
              % label, "keyUsage is %r" % ku)

    subject = text(pem, "-subject")
    check(host in subject, "%s: subject CN is the requested host" % label,
          subject.strip())

    san = ext(pem, "subjectAltName")
    check(san is not None and host in san,
          "%s: requested host present as a DNS subjectAltName" % label,
          "subjectAltName is %r" % san)

    pubkey = text(pem, "-text")
    check("rsaEncryption" in pubkey or "RSA Public-Key" in pubkey
          or "Public Key Algorithm: rsaEncryption" in pubkey,
          "%s: generated cert carries squid's RSA key" % label)


def check_database(workdir):
    db = os.path.join(workdir, "ssl_db")
    run(CERTGEN, "-c", "-s", db, "-M", "4MB")

    host = "db.example.com"
    origin = make_origin(workdir, host, "ec")

    gen = CertGen(db_path=db)
    first = gen.request(host, origin)
    second = gen.request(host, origin)
    gen.close()

    check(first == second,
          "db: a repeated request is served from the database",
          "certgen minted a second certificate, so the database key derived "
          "from the mimicked cert's signature did not match")

    index = os.path.join(db, "index.txt")
    check(os.path.exists(index), "db: index.txt was written")
    if not os.path.exists(index):
        return

    with open(index) as f:
        rows = [ln.rstrip("\n").split("\t") for ln in f if ln.strip()]
    check(len(rows) == 1, "db: exactly one entry was stored",
          "index.txt has %d rows" % len(rows))
    if not rows:
        return

    # Columns: key, expiry date, revocation date, serial, filename, subject.
    exp_date = rows[0][1]
    check(re.fullmatch(r"\d{12,15}Z", exp_date) is not None,
          "db: expiry column holds a well-formed ASN.1 time",
          "got %r; certificate_db.cc reads the raw ASN1_STRING bytes, so a "
          "truncated or empty value means ASN1_STRING_get0_data()/"
          "ASN1_STRING_length() returned nothing" % exp_date)

    not_after = text(first, "-enddate").strip()
    check(not_after.startswith("notAfter="),
          "db: generated certificate has a readable notAfter")

    subject_col = rows[0][-1]
    check(host in subject_col,
          "db: subject column recorded by OneLineSummary()",
          "got %r" % subject_col)


def main():
    global CA_CRT, CA_KEY

    if not os.path.exists(CERTGEN):
        print("SKIP: %s is not installed (squid-openssl missing)" % CERTGEN)
        return 77
    if not shutil.which("openssl"):
        print("SKIP: the openssl command is not available")
        return 77

    workdir = tempfile.mkdtemp(
        prefix="squid-certmimic-", dir=os.environ.get("AUTOPKGTEST_TMP"))
    try:
        CA_CRT, CA_KEY = make_ca(workdir)
        print("openssl: %s" % run("openssl", "version").strip())

        for keytype, expect_ke in (("ec", True), ("rsa", False)):
            host = "%s-origin.example.com" % keytype
            origin = make_origin(workdir, host, keytype)
            gen = CertGen()
            pem = gen.request(host, origin)
            gen.close()
            check_generated(pem, host, expect_ke, keytype)

        check_database(workdir)
    finally:
        if not os.environ.get("AUTOPKGTEST_TMP"):
            shutil.rmtree(workdir, ignore_errors=True)

    print("\n%d checks, %d failures" % (checks, len(failures)))
    for f in failures:
        print("  failed: %s" % f)
    return 1 if failures else 0


if __name__ == "__main__":
    sys.exit(main())
