"""
adcs_cert_request.py
--------------------
Génère une clé privée + une CSR (EKU Client Authentication + Server Authentication),
la soumet à une autorité de certification Microsoft (AD CS), récupère le certificat
émis et produit un fichier PFX prêt à importer.

Deux modes de soumission :
  * certreq : utilise certreq.exe (Windows, machine du domaine, auth Kerberos implicite)
  * web     : utilise l'interface Web Enrollment (/certsrv) en NTLM
              -> fonctionne aussi depuis une machine hors domaine (workgroup)

Dépendances :
  pip install cryptography requests requests-ntlm   # requests* seulement pour le mode web

Exemples :
  # Machine du domaine, via certreq
  python adcs_cert_request.py --cn srv01.corp.local --san srv01 --san srv01.corp.local \
      --template ClientServerAuth --method certreq --ca-config "ca01.corp.local\\CORP-CA"

  # Machine en workgroup, via Web Enrollment
  python adcs_cert_request.py --cn gw01.dmz.local --template ClientServerAuth \
      --method web --ca-host ca01.corp.local --user "CORP\\svc_enroll"
"""

import argparse
import getpass
import os
import re
import subprocess
import sys
import tempfile
from pathlib import Path

from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.serialization import pkcs12
from cryptography.x509.oid import ExtendedKeyUsageOID, NameOID


# --------------------------------------------------------------------------- #
# 1. Clé + CSR
# --------------------------------------------------------------------------- #
def build_csr(cn, sans, key_size, org=None, ou=None, country=None):
    key = rsa.generate_private_key(public_exponent=65537, key_size=key_size)

    name_attrs = [x509.NameAttribute(NameOID.COMMON_NAME, cn)]
    if ou:
        name_attrs.append(x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, ou))
    if org:
        name_attrs.append(x509.NameAttribute(NameOID.ORGANIZATION_NAME, org))
    if country:
        name_attrs.append(x509.NameAttribute(NameOID.COUNTRY_NAME, country))

    # Le CN est toujours repris dans le SAN (les clients TLS modernes ignorent le CN)
    dns_names = list(dict.fromkeys([cn] + (sans or [])))

    csr = (
        x509.CertificateSigningRequestBuilder()
        .subject_name(x509.Name(name_attrs))
        .add_extension(
            x509.SubjectAlternativeName([x509.DNSName(n) for n in dns_names]),
            critical=False,
        )
        .add_extension(
            x509.ExtendedKeyUsage(
                [ExtendedKeyUsageOID.SERVER_AUTH, ExtendedKeyUsageOID.CLIENT_AUTH]
            ),
            critical=False,
        )
        .add_extension(
            x509.KeyUsage(
                digital_signature=True, key_encipherment=True,
                content_commitment=False, data_encipherment=False,
                key_agreement=False, key_cert_sign=False, crl_sign=False,
                encipher_only=False, decipher_only=False,
            ),
            critical=True,
        )
        .sign(key, hashes.SHA256())
    )
    return key, csr


# --------------------------------------------------------------------------- #
# 2a. Soumission via certreq.exe
# --------------------------------------------------------------------------- #
def submit_certreq(csr_path, template, ca_config, out_dir):
    if os.name != "nt":
        sys.exit("Le mode certreq nécessite Windows (certreq.exe).")
    cer_path = out_dir / (csr_path.stem + ".cer")
    cmd = ["certreq.exe", "-submit", "-q",
           "-attrib", f"CertificateTemplate:{template}"]
    if ca_config:
        cmd += ["-config", ca_config]
    cmd += [str(csr_path), str(cer_path)]

    print(f"[*] {' '.join(cmd)}")
    res = subprocess.run(cmd, capture_output=True, text=True)
    print(res.stdout.strip())
    if res.returncode != 0 or not cer_path.exists():
        req_id = re.search(r"RequestId:\s*\"?(\d+)", res.stdout)
        if req_id:
            sys.exit(f"[!] Requête en attente ou refusée (RequestId {req_id.group(1)}). "
                     f"Récupération ultérieure : certreq -retrieve {req_id.group(1)} out.cer")
        sys.exit(f"[!] Échec certreq : {res.stderr.strip()}")
    return x509.load_pem_x509_certificate(cer_path.read_bytes()) \
        if b"BEGIN CERTIFICATE" in cer_path.read_bytes() \
        else x509.load_der_x509_certificate(cer_path.read_bytes())


# --------------------------------------------------------------------------- #
# 2b. Soumission via Web Enrollment (/certsrv)
# --------------------------------------------------------------------------- #
def submit_web(csr_pem, template, ca_host, user, password, verify):
    import requests
    from requests_ntlm import HttpNtlmAuth

    base = f"https://{ca_host}/certsrv"
    s = requests.Session()
    s.auth = HttpNtlmAuth(user, password)
    s.verify = verify

    data = {
        "Mode": "newreq",
        "CertRequest": csr_pem,
        "CertAttrib": f"CertificateTemplate:{template}",
        "TargetStoreFlags": "0",
        "SaveCert": "yes",
    }
    print(f"[*] POST {base}/certfnsh.asp (modèle {template})")
    r = s.post(f"{base}/certfnsh.asp", data=data, timeout=60)
    if r.status_code == 401:
        sys.exit("[!] Authentification refusée (vérifier compte / NTLM activé sur IIS).")
    r.raise_for_status()

    m = re.search(r"certnew\.cer\?ReqID=(\d+)", r.text)
    if not m:
        pending = re.search(r"Your Request Id is (\d+)", r.text)
        if pending:
            sys.exit(f"[!] Requête en attente d'approbation (ReqID {pending.group(1)}). "
                     f"Récupérer plus tard : {base}/certnew.cer?ReqID={pending.group(1)}&Enc=b64")
        err = re.search(r"The disposition message is \"([^\"]+)", r.text)
        sys.exit(f"[!] Requête refusée : {err.group(1) if err else 'réponse inattendue de /certsrv'}")

    req_id = m.group(1)
    print(f"[+] Certificat émis (ReqID {req_id}), téléchargement…")
    c = s.get(f"{base}/certnew.cer", params={"ReqID": req_id, "Enc": "b64"}, timeout=60)
    c.raise_for_status()
    return x509.load_pem_x509_certificate(c.content)


# --------------------------------------------------------------------------- #
# 3. Main
# --------------------------------------------------------------------------- #
def main():
    p = argparse.ArgumentParser(description="CSR client/serveur + soumission AD CS")
    p.add_argument("--cn", required=True, help="Common Name (FQDN de la machine)")
    p.add_argument("--san", action="append", default=[], help="SAN DNS (répétable)")
    p.add_argument("--org"); p.add_argument("--ou"); p.add_argument("--country")
    p.add_argument("--key-size", type=int, default=2048, choices=[2048, 3072, 4096])
    p.add_argument("--template", required=True, help="Nom (pas le display name) du modèle AD CS")
    p.add_argument("--method", choices=["certreq", "web"], default="certreq")
    p.add_argument("--ca-config", help='certreq : "serveurCA\\NomCA" (sinon fenêtre de choix)')
    p.add_argument("--ca-host", help="web : FQDN du serveur Web Enrollment")
    p.add_argument("--user", help='web : compte "DOMAINE\\user"')
    p.add_argument("--ca-bundle", help="web : fichier PEM de la CA racine pour valider le TLS")
    p.add_argument("--insecure", action="store_true", help="web : ne pas valider le TLS (test)")
    p.add_argument("--out-dir", default=".", help="Répertoire de sortie")
    p.add_argument("--no-pfx", action="store_true", help="Ne pas générer de PFX")
    a = p.parse_args()

    out = Path(a.out_dir); out.mkdir(parents=True, exist_ok=True)
    base = a.cn.replace("*", "wildcard")

    # Clé + CSR
    key, csr = build_csr(a.cn, a.san, a.key_size, a.org, a.ou, a.country)
    pfx_pwd = None if a.no_pfx else getpass.getpass("Mot de passe du PFX (et de la clé) : ")
    enc = (serialization.BestAvailableEncryption(pfx_pwd.encode())
           if pfx_pwd else serialization.NoEncryption())

    key_path = out / f"{base}.key"
    key_path.write_bytes(key.private_bytes(serialization.Encoding.PEM,
                                           serialization.PrivateFormat.PKCS8, enc))
    try:
        os.chmod(key_path, 0o600)
    except OSError:
        pass
    csr_pem = csr.public_bytes(serialization.Encoding.PEM).decode()
    csr_path = out / f"{base}.csr"
    csr_path.write_text(csr_pem)
    print(f"[+] Clé : {key_path}\n[+] CSR : {csr_path}")

    # Soumission
    if a.method == "certreq":
        cert = submit_certreq(csr_path, a.template, a.ca_config, out)
    else:
        if not (a.ca_host and a.user):
            sys.exit("--ca-host et --user sont requis en mode web.")
        pwd = os.environ.get("ADCS_PASSWORD") or getpass.getpass(f"Mot de passe de {a.user} : ")
        verify = False if a.insecure else (a.ca_bundle or True)
        cert = submit_web(csr_pem, a.template, a.ca_host, a.user, pwd, verify)

    cer_path = out / f"{base}.cer"
    cer_path.write_bytes(cert.public_bytes(serialization.Encoding.PEM))
    print(f"[+] Certificat : {cer_path}")
    print(f"    Sujet  : {cert.subject.rfc4514_string()}")
    print(f"    Émis par : {cert.issuer.rfc4514_string()}")
    print(f"    Valide jusqu'au : {cert.not_valid_after_utc:%Y-%m-%d}")

    # Contrôle des EKU réellement émis (le modèle de la CA a le dernier mot)
    try:
        ekus = cert.extensions.get_extension_for_class(x509.ExtendedKeyUsage).value
        for oid, label in [(ExtendedKeyUsageOID.SERVER_AUTH, "Server Auth"),
                           (ExtendedKeyUsageOID.CLIENT_AUTH, "Client Auth")]:
            print(f"    {label:12}: {'OK' if oid in ekus else 'ABSENT !'}")
    except x509.ExtensionNotFound:
        print("    [!] Aucun EKU dans le certificat émis.")

    # PFX
    if not a.no_pfx:
        pfx = pkcs12.serialize_key_and_certificates(
            name=a.cn.encode(), key=key, cert=cert, cas=None,
            encryption_algorithm=serialization.BestAvailableEncryption(pfx_pwd.encode()),
        )
        pfx_path = out / f"{base}.pfx"
        pfx_path.write_bytes(pfx)
        print(f"[+] PFX : {pfx_path}")
        print(f'    Import (Windows, admin) : certutil -f -p "<mdp>" -importpfx "{pfx_path}"')


if __name__ == "__main__":
    main()