#!/usr/bin/python3
import argparse
import json
import os
import sys

parser = argparse.ArgumentParser(description="Create the credential configuration of the Storage Transfer agent container.")
parser.add_argument("--project-number", required=True, help="number of the project of the workload identity pool")
parser.add_argument("--pool", required=True, help="workload identity pool ID")
parser.add_argument("--provider", required=True, help="X.509 provider ID")
parser.add_argument("--service-account", required=True, help="email of the impersonated service account")
parser.add_argument("--output", default="/etc/sts-agent", help="credential directory of the agent")
args = parser.parse_args()

output = os.path.abspath(args.output)
cert = os.path.join(output, "client_cert.pem")
key = os.path.join(output, "client_key.pem")

# The client certificate and key must be installed in the credential directory first.
for path in (cert, key):
    if not os.path.isfile(path):
        sys.exit(f"{path} does not exist")

# The credential directory holds the private key and is mounted read only in the agent container.
os.chmod(output, 0o750)

cred = {
    "universe_domain": "googleapis.com",
    "type": "external_account",
    "audience": f"//iam.googleapis.com/projects/{args.project_number}/locations/global/workloadIdentityPools/{args.pool}/providers/{args.provider}",
    "subject_token_type": "urn:ietf:params:oauth:token-type:mtls",
    "token_url": "https://sts.mtls.googleapis.com/v1/token",
    "credential_source": {
        "certificate": {
            "use_default_certificate_config": True,
        }
    },
    "service_account_impersonation_url": f"https://iamcredentials.mtls.googleapis.com/v1/projects/-/serviceAccounts/{args.service_account}:generateAccessToken",
}

# The container mounts the credential directory at the same path as on the host.
certificate_config = {
    "cert_configs": {
        "workload": {
            "cert_path": cert,
            "key_path": key,
        }
    }
}

for name, content in (("cred.json", cred), ("certificate_config.json", certificate_config)):
    with open(os.path.join(output, name), "w") as file:
        json.dump(content, file, indent=2)
        file.write("\n")