psql

#!/usr/bin/env python3

import argparse
import base64
import hashlib
import json
import shlex
import signal
import subprocess
import sys
from webbrowser import get


def _get_secret(keyvault: str, keyvault_key: str) -> str:
    json_str=  subprocess.run([
                            "az",
                    "keyvault",
                    "secret",
                    "show",
                    f"--vault-name={keyvault}",
                    f"--name={keyvault_key}",
    ], stdout=subprocess.PIPE, encoding="utf-8", check=True).stdout
    return json.loads(json_str)["value"]

def _get_secrets(keyvault: str, database: str) -> dict[str, str]:
    return {
        "database": _get_secret(keyvault, f"pg-{database}-database"),
        "username": _get_secret(keyvault, f"pg-{database}-username"),
        "password": _get_secret(keyvault, f"pg-{database}-password"),
    }


def _get_env(decoded_secrets: dict[str, str]) -> dict[str, str]:
    return {
        "PGHOST": "odyssey.apps.gs-ch-prod.camptocamp.com",
        "PGPORT": "5432",
        "PGDATABASE": decoded_secrets["database"],
        "PGUSER": decoded_secrets["username"],
        "PGPASSWORD": decoded_secrets["password"],
        "PGSSLMODE": "require",
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run psql with the credentials from Azure Keyvault.",
                                     usage= "scripts/psql gmf-demoprod-kv gmf")
    parser.add_argument("--odyssey", action="store_true", help="Get the configuration for Odyssey")
    parser.add_argument("--display", action="store_true", help="Display the credentials")
    parser.add_argument(
        "--dry-run", action="store_true", help="Display the command to run without executing it"
    )
    parser.add_argument("keyvault", nargs=1, help="The keyvault name")
    # optional argument
    parser.add_argument("database", nargs='?', help="The database nick name")
    parser.add_argument("command", nargs=argparse.REMAINDER)
    args = parser.parse_args()

    keyvault = args.keyvault[0]
    if args.database is None:
        proc = subprocess.run(["az", "keyvault", "secret", "list", f"--vault-name={keyvault}"], check=True, stdout=subprocess.PIPE)
        print("Available databases:")
        for secret in json.loads(proc.stdout):
            if secret["name"].startswith("pg-") and secret["name"].endswith("-database"):
                print(secret["name"][len("pg-"):-len("-database")])
        sys.exit(0)
    database = args.database

    decoded_secrets = _get_secrets(keyvault, database)

    if args.odyssey:
        short_name = keyvault[:-6] if keyvault.endswith("int-kv") else keyvault[:-7]
        env = "int" if keyvault.endswith("int-kv") else "prod"
        print(
            f"""
To add in this file https://github.com/camptocamp/argocd-shelter-cluster-apps/blob/main/apps/prod/gs/ch/odyssey/secrets.yaml
With SOPS:

            {short_name}-{database}-{env}:
                storage: gmf_{env}
                username: {decoded_secrets["username"]}
                database: {decoded_secrets["database"]}
                password: md5{hashlib.md5((decoded_secrets['password'] + decoded_secrets['username']).encode()).hexdigest()}
                options:
                    authentication: md5
"""
        )
        sys.exit(0)

    env = _get_env(decoded_secrets)
    if args.display:
        for key, value in env.items():
            print(f"{key}={value}")
        sys.exit(0)
    if args.dry_run:
        print(" ".join(["=".join(e) for e in env.items()]) + " " + shlex.join(args.command or ["psql"]))
        sys.exit(0)

    print(f'Used database: {decoded_secrets["database"]}')

    signal.signal(signal.SIGINT, signal.SIG_IGN)
    proc = subprocess.run(args.command or ["psql"], env=env)  # pylint: disable=subprocess-run-check
    sys.exit(proc.returncode)