#!/usr/bin/env python3

import argparse
import base64
import pathlib
import tempfile
import uuid

from jwcrypto import jwk

def create_jwk(output_dir):
    key_verify = jwk.JWK.generate(
        kty="RSA",
        size=4096,
        alg="PS512",
        kid=str(uuid.uuid4()),
        key_ops=["verify"],
    )

    key_wrapkey = jwk.JWK.generate(
        kty="RSA",
        size=4096,
        alg="RSA-OAEP-256",
        kid=str(uuid.uuid4()),
        key_ops=["wrapKey"],
    )

    #aus dem Schlüsselpaar exportierte Public und Private Keys als String
    publicKey_verify_str = key_verify.export(private_key=False)
    publicKey_wrapkey_str = key_wrapkey.export(private_key=False)
    
    privateKey_sign_str = key_verify.export_private(as_dict=False)
    privateKey_unwrapkey_str = key_wrapkey.export_private(as_dict=False)

    #aus dem Schlüsselpaar exportierte Public und Private Keys als JSON
    publicKey_verify = jwk.JWK.from_json(publicKey_verify_str)
    publicKey_wrapkey = jwk.JWK.from_json(publicKey_wrapkey_str)

    privateKey_sign = jwk.JWK.from_json(privateKey_sign_str)
    privateKey_unwrapkey = jwk.JWK.from_json(privateKey_unwrapkey_str)

    #ergänze Public Key JSON um x5c Element
    # Auskommentiert, da dem Schlüssel aktuell kein Zertifikat zugrunde liegt und der Dummy Wert fehlerhaft ist bzw. bei der Verarbeitung nicht akzeptiert wird
    # publicKey_verify["x5c"] = [base64.urlsafe_b64encode('dummy'.encode('UTF-8')).decode('utf-8').rstrip("=")]
    # publicKey_wrapkey["x5c"] = [base64.urlsafe_b64encode('dummy'.encode('UTF-8')).decode('utf-8').rstrip("=")]

    output_dir.mkdir(parents=True, exist_ok=True)

    keySet_file = pathlib.Path(output_dir, "set-public-keys.json")
    publicKey_wrapkey_file = pathlib.Path(output_dir, "publicKey_encryption.json")
    publicKey_verify_file = pathlib.Path(output_dir, "publicKey_signature_verification.json")
    privateKey_unwrapkey_file = pathlib.Path(output_dir, "privateKey_decryption.json")
    privateKey_sign_file = pathlib.Path(output_dir, "privateKey_signing.json")

    jwks = jwk.JWKSet()
    jwks.add(publicKey_verify)
    jwks.add(publicKey_wrapkey)

    with open(keySet_file, "wb") as f:
        exp = jwks.export()
        f.write(exp.encode("UTF-8"))

    with open(publicKey_verify_file, "wb") as f:
        exp = publicKey_verify.export()
        f.write(exp.encode("UTF-8"))

    with open(publicKey_wrapkey_file, "wb") as f:
        exp = publicKey_wrapkey.export()
        f.write(exp.encode("UTF-8"))

    with open(privateKey_sign_file, "wb") as f:
        exp = privateKey_sign.export()
        f.write(exp.encode("UTF-8"))

    with open(privateKey_unwrapkey_file, "wb") as f:
        exp = privateKey_unwrapkey.export()
        f.write(exp.encode("UTF-8"))

    print(f"🔒 Wrote JWK representation of encryption public key (key_use=wrapKey) to {publicKey_wrapkey_file}")
    print(f"🔒 Wrote JWK representation of signature validation public key (key_use=verify) to {publicKey_verify_file}")
    print("Please upload these keys when creating a destination in the self service portal.")
    print()
    print(f"🔒 Wrote JWKS of Public Keys to {keySet_file}")
    print("This key set can be used to update (rotate) keys via the Submission-API (PUT /destinations/\{destinationID\})")
    print()
    print(f"🔒 Wrote JWK representation of decryption private key (key_use=unwrapKey) to {privateKey_unwrapkey_file}")
    print(f"🔒 Wrote JWK representation of signing private key (key_use=sign) to {privateKey_sign_file}")
    print("These keys can be used to sign and decrypt in your client application.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate SET JWKS.")
    parser.add_argument(
        "-o",
        "--output",
        default=tempfile.mkdtemp(),
        help="Directory to store the generated SET JWKS in. Default: a temporary directory",
        type=pathlib.Path,
    )
    args = parser.parse_args()

    create_jwk(args.output)