import subprocess, datetime, os, hashlib, ipaddress, re
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.x509.oid import NameOID
from config import *

TMP_DIR = "/var/lib/certauth/tmp"
os.makedirs(TMP_DIR, exist_ok=True)

def get_root_pub_key():
    with open(YK_PUB_ROOT, "rb") as f:
        return serialization.load_pem_public_key(f.read())

def get_int_pub_key():
    with open(YK_PUB_INT, "rb") as f:
        return serialization.load_pem_public_key(f.read())

def get_root_ca_cert():
    with open(ROOT_CA_PATH, "rb") as f:
        return x509.load_pem_x509_certificate(f.read())

def der_len(n):
    if n < 0x80: return bytes([n])
    elif n < 0x100: return bytes([0x81, n])
    return bytes([0x82, n>>8, n&0xff])

def sign_tbs_with_yk(tbs_bytes, yk_pin, token_label="certauth Intermediate CA"):
    tbs_file = os.path.join(TMP_DIR, "tbs_sign.der")
    sig_file = os.path.join(TMP_DIR, "sig_out.bin")
    with open(tbs_file, "wb") as f:
        f.write(tbs_bytes)
    r = subprocess.run([
        "sudo", "pkcs11-tool", "--module", PKCS11_MODULE,
        "--login", "--pin", yk_pin,
        "--sign", "--mechanism", "ECDSA-SHA384",
        "--token-label", token_label,
        "--label", "SIGN key",
        "--input-file", tbs_file,
        "--output-file", sig_file
    ], capture_output=True, text=True)
    if r.returncode != 0:
        return None, r.stderr
    with open(sig_file, "rb") as f:
        raw = f.read()
    rb = raw[:48].lstrip(b"\x00") or b"\x00"
    sb = raw[48:].lstrip(b"\x00") or b"\x00"
    if rb[0] & 0x80: rb = b"\x00" + rb
    if sb[0] & 0x80: sb = b"\x00" + sb
    r_der = b"\x02" + bytes([len(rb)]) + rb
    s_der = b"\x02" + bytes([len(sb)]) + sb
    seq = b"\x30" + bytes([len(r_der+s_der)]) + r_der + s_der
    bs = b"\x00" + seq
    return b"\x03" + bytes([len(bs)]) + bs, None

def build_leaf_cert(cn, sans, days=365):
    root_cert = get_root_ca_cert()
    int_pub = get_int_pub_key()
    root_pub = get_root_pub_key()
    leaf_key = ec.generate_private_key(ec.SECP384R1())
    subject = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Home"),
        x509.NameAttribute(NameOID.COMMON_NAME, cn),
    ])
    issuer = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, "US"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "Home"),
        x509.NameAttribute(NameOID.COMMON_NAME, "certauth Intermediate CA"),
    ])
    builder = (x509.CertificateBuilder()
        .subject_name(subject).issuer_name(issuer)
        .public_key(leaf_key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(datetime.datetime.now(datetime.timezone.utc))
        .not_valid_after(datetime.datetime.now(datetime.timezone.utc) + datetime.timedelta(days=days))
        .add_extension(x509.BasicConstraints(ca=False, path_length=None), critical=True)
        .add_extension(x509.KeyUsage(
            digital_signature=True, key_encipherment=True,
            key_cert_sign=False, crl_sign=False,
            content_commitment=False, data_encipherment=False,
            key_agreement=False, encipher_only=False, decipher_only=False), critical=True)
        .add_extension(x509.ExtendedKeyUsage([
            x509.oid.ExtendedKeyUsageOID.SERVER_AUTH,
            x509.oid.ExtendedKeyUsageOID.CLIENT_AUTH,
        ]), critical=False)
        .add_extension(x509.SubjectKeyIdentifier.from_public_key(leaf_key.public_key()), critical=False)
        .add_extension(x509.AuthorityKeyIdentifier.from_issuer_public_key(int_pub), critical=False))
    if sans:
        san_list = []
        for s in sans.split(","):
            s = s.strip()
            # Check if it's an IP address
            if re.match(r"^\d{1,3}(\.\d{1,3}){3}$", s):
                san_list.append(x509.IPAddress(ipaddress.ip_address(s)))
            else:
                san_list.append(x509.DNSName(s))
        builder = builder.add_extension(x509.SubjectAlternativeName(san_list), critical=False)
    tmp = ec.generate_private_key(ec.SECP384R1())
    temp = builder.sign(tmp, hashes.SHA384())
    td = temp.public_bytes(serialization.Encoding.DER)
    o = 1
    if td[o] & 0x80: n = td[o] & 0x7f; o += 1 + n
    else: o += 1
    tbs_start = o
    o += 1
    if td[o] & 0x80: n = td[o] & 0x7f; tl = int.from_bytes(td[o+1:o+1+n], "big"); o += 1 + n
    else: tl = td[o]; o += 1
    tbs_end = o + tl
    tbs_full = td[tbs_start:tbs_end]
    alg_start = tbs_end
    o2 = alg_start + 1
    if td[o2] & 0x80: n = td[o2] & 0x7f; al = int.from_bytes(td[o2+1:o2+1+n], "big"); o2 += 1 + n
    else: al = td[o2]; o2 += 1
    alg_full = td[alg_start:o2+al]
    new_sig, err = sign_tbs_with_yk(tbs_full, YK_INT_PIN)
    if new_sig is None: return None, err
    content = tbs_full + alg_full + new_sig
    cl = len(content)
    final = b"\x30\x82" + bytes([cl>>8, cl&0xff]) + content
    der_file = os.path.join(TMP_DIR, "leaf.der")
    pem_file = os.path.join(TMP_DIR, "leaf.pem")
    with open(der_file, "wb") as f: f.write(final)
    r = subprocess.run(["openssl", "x509", "-inform", "DER", "-outform", "PEM",
                         "-in", der_file, "-out", pem_file],
                       capture_output=True, text=True)
    if r.returncode != 0: return None, r.stderr
    with open(pem_file) as f: leaf_pem = f.read()
    key_pem = leaf_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption()
    ).decode()
    cert = x509.load_pem_x509_certificate(leaf_pem.encode())
    serial = format(cert.serial_number, 'x')
    return {"cert_pem": leaf_pem, "key_pem": key_pem, "serial": serial,
            "expires_at": cert.not_valid_after.isoformat()}, None