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

parser = argparse.ArgumentParser(description="Renew the client certificate from Vault with AppRole.")
parser.add_argument(
    "--unwrap",
    action="store_true",
    help="read a response wrapping token and store the unwrapped secret ID",
)
parser.add_argument(
    "--login-test",
    action="store_true",
    help="log in with AppRole, print the token policies and lifetime, and exit",
)
args = parser.parse_args()


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


vault_addr = setting("VAULT_ADDR").rstrip("/")
context = ssl.create_default_context(cafile=os.environ.get("VAULT_CACERT") or None)


def api(path, body=None, 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 or {}).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()}")


role_id_file = setting("ROLE_ID_FILE", "/etc/vault-cert/role_id")
secret_id_file = setting("SECRET_ID_FILE", "/etc/vault-cert/secret_id")


def login():
    with open(role_id_file) as file:
        role_id = file.read().strip()
    with open(secret_id_file) as file:
        secret_id = file.read().strip()
    return api("auth/approle/login", {"role_id": role_id, "secret_id": secret_id})["auth"]


# Store the secret ID that is wrapped in a single use response wrapping token.
if args.unwrap:
    wrapping_token = getpass.getpass("Wrapping token: ")
    secret_id = api("sys/wrapping/unwrap", token=wrapping_token)["data"]["secret_id"]
    descriptor = os.open(secret_id_file, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    os.fchmod(descriptor, 0o600)
    with os.fdopen(descriptor, "w") as file:
        file.write(secret_id)
    print(f"The secret ID was written to {secret_id_file}")
    sys.exit(0)

# Check the AppRole credentials without issuing a certificate.
if args.login_test:
    auth = login()
    print(f"Login succeeded: token type {auth['token_type']}, policies {auth['token_policies']}, valid for {auth['lease_duration']} seconds")
    sys.exit(0)

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 = os.environ.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,
    )
    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")
    token = login()["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)