#!/usr/bin/python3
import argparse
import json
import os
import ssl
import subprocess
import sys
import tempfile
import urllib.error
import urllib.request

parser = argparse.ArgumentParser(description="Renew a certificate from Vault with AppRole.")
parser.add_argument("config", help="path to the JSON configuration file of the instance")
args = parser.parse_args()

with open(args.config) as file:
    config = json.load(file)


def setting(name, default=None):
    value = config.get(name, default)
    if value is None:
        sys.exit(f"{name} is not set in {args.config}")
    return value


vault_addr = setting("vault_addr").rstrip("/")
context = ssl.create_default_context(cafile=config.get("vault_cacert") or None)


def api(path, body, token=None):
    headers = {"Content-Type": "application/json"}
    if token:
        headers["X-Vault-Token"] = token
    request = urllib.request.Request(
        f"{vault_addr}/v1/{path}",
        data=json.dumps(body).encode(),
        headers=headers,
        method="POST",
    )
    try:
        with urllib.request.urlopen(request, context=context) as response:
            return json.load(response)
    except urllib.error.HTTPError as error:
        sys.exit(f"Vault returned {error.code} for {path}: {error.read().decode()}")


cert_dir = setting("cert_dir", "/var/lib/vault-cert")
common_name = setting("cert_cn")
sign_path = setting("sign_path", "pki_int/sign/wif-client")
renew_before_days = int(setting("renew_before_days", 30))
cert_ttl = setting("cert_ttl", "8760h")
hook = config.get("hook", "")

cert = os.path.join(cert_dir, "client_cert.pem")
key = os.path.join(cert_dir, "client_key.pem")

# Do nothing while the certificate is valid for longer than the renewal window.
if os.path.isfile(cert) and os.path.isfile(key):
    valid = subprocess.run(
        ["openssl", "x509", "-checkend", str(renew_before_days * 86400), "-noout", "-in", cert],
        capture_output=True,
        check=False,
    )
    if valid.returncode == 0:
        print(f"The certificate is valid for more than {renew_before_days} days, nothing to do")
        sys.exit(0)

os.makedirs(cert_dir, mode=0o750, exist_ok=True)
with tempfile.TemporaryDirectory(dir=cert_dir, prefix=".renew.") as workdir:
    new_key = os.path.join(workdir, "client_key.pem")
    new_cert = os.path.join(workdir, "client_cert.pem")
    csr_file = os.path.join(workdir, "client.csr")

    print("Logging in to Vault with AppRole")
    login = {"role_id": setting("role_id"), "secret_id": setting("secret_id")}
    token = api("auth/approle/login", login)["auth"]["client_token"]

    print(f"Generating the key and the CSR for {common_name}")
    subprocess.run(
        ["openssl", "genpkey", "-algorithm", "RSA", "-pkeyopt", "rsa_keygen_bits:4096", "-out", new_key],
        check=True,
    )
    subprocess.run(
        ["openssl", "req", "-new", "-key", new_key, "-subj", f"/CN={common_name}", "-out", csr_file],
        check=True,
    )

    print("Signing the CSR")
    with open(csr_file) as file:
        csr = file.read()
    signed = api(sign_path, {"csr": csr, "common_name": common_name, "ttl": cert_ttl}, token)
    with open(new_cert, "w") as file:
        file.write(signed["data"]["certificate"].rstrip("\n") + "\n")

    # Replace the files only when the new certificate matches the new key.
    cert_public = subprocess.run(
        ["openssl", "x509", "-in", new_cert, "-noout", "-pubkey"],
        capture_output=True, text=True, check=True,
    ).stdout
    key_public = subprocess.run(
        ["openssl", "pkey", "-in", new_key, "-pubout"],
        capture_output=True, text=True, check=True,
    ).stdout
    if cert_public != key_public:
        sys.exit("The issued certificate does not match the generated key")

    os.chmod(new_key, 0o600)
    os.chmod(new_cert, 0o644)
    os.replace(new_key, key)
    os.replace(new_cert, cert)

end_date = subprocess.run(
    ["openssl", "x509", "-in", cert, "-noout", "-enddate"],
    capture_output=True, text=True, check=True,
).stdout.strip().split("=", 1)[1]
print(f"The certificate was renewed, it expires on {end_date}")

if hook:
    if os.access(hook, os.X_OK) and os.path.isfile(hook):
        print(f"Running the hook {hook}")
        subprocess.run([hook], env={**os.environ, "CERT_FILE": cert, "KEY_FILE": key}, check=True)
    else:
        print(f"The hook {hook} is not an executable file, skipping it", file=sys.stderr)