hashicorp-kmip-setup.py

#!/usr/bin/env python3
"""
HashiCorp Vault Enterprise KMIP Setup
"""

import argparse
import json
import os
import shutil
import subprocess
import sys
import time
from pathlib import Path


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def log(level: str, msg: str) -> None:
    print(f"[{level}] {msg}", flush=True)


def info(msg: str) -> None:
    log("INFO", msg)


def error(msg: str) -> None:
    log("ERROR", msg)


def run(cmd: list[str], capture_output: bool = False, check: bool = True) -> subprocess.CompletedProcess:
    """Run a command, optionally capturing output."""
    return subprocess.run(cmd, capture_output=capture_output, text=True, check=check)


def docker_exec(container: str, shell_cmd: str, capture_output: bool = False) -> subprocess.CompletedProcess:
    """Run a shell command inside a Docker container."""
    return run(
        ["docker", "exec", container, "sh", "-c", shell_cmd],
        capture_output=capture_output,
    )


# ---------------------------------------------------------------------------
# Steps
# ---------------------------------------------------------------------------

def setup_directories(config_dir: Path, data_dir: Path, log_dir: Path, certs_dir: Path) -> None:
    for d in (config_dir, data_dir, log_dir, certs_dir):
        d.mkdir(parents=True, exist_ok=True)
        run(["sudo", "chown", "-R", "100:1000", str(d)])

    # certs dir needs to be writable by the host user as well
    run(["sudo", "chmod", "-R", "777", str(certs_dir)])


def check_license(license_path: str | None) -> str:
    info("Checking for license file...")

    # First check environment variable (like bash: export VAULT_LICENSE=...)
    env_license = os.environ.get("VAULT_LICENSE")
    if env_license:
        info("Using license from environment variable VAULT_LICENSE")
        return env_license.strip()

    # Fall back to --license file path
    if not license_path:
        error("No license provided. Either set VAULT_LICENSE env var or use --license=/path/to/vault.hclic")
        sys.exit(1)

    p = Path(license_path)
    if not p.is_file():
        error(f"License file not found at: {license_path}")
        info("Please ensure the license file path is correct")
        sys.exit(1)

    info(f"Using license from: {license_path}")
    return p.read_text().strip()


def create_vault_hcl(vault_hcl: Path) -> None:
    if vault_hcl.is_file():
        info(f"Vault HCL config already exists at {vault_hcl}")
        return

    info(f"Creating Vault HCL config at {vault_hcl}")
    vault_hcl.write_text(
        '# vault.hcl autogenerated\n'
        'storage "raft" {\n'
        '  path    = "/vault/data"\n'
        '  node_id = "node1"\n'
        '}\n\n'
        'listener "tcp" {\n'
        '  address     = "0.0.0.0:8200"\n'
        '  tls_disable = 1\n'
        '}\n\n'
        'api_addr      = "http://127.0.0.1:8200"\n'
        'cluster_addr  = "http://127.0.0.1:8201"\n'
        'disable_mlock = true\n'
        'log_level     = "trace"\n'
    )


def start_vault_container(
    container_name: str,
    certs_dir: Path,
    vault_hcl: Path,
    data_dir: Path,
    log_dir: Path,
    license_content: str,
) -> None:
    # Check current container state
    result = run(
        ["docker", "inspect", "-f", "{{.State.Status}}", container_name],
        capture_output=True,
        check=False,
    )
    status = result.stdout.strip() if result.returncode == 0 else ""

    if status == "running":
        info("HashiCorp Vault container already running")
        return

    if status == "exited":
        info("Starting existing HashiCorp Vault container")
        run(["docker", "start", container_name], capture_output=True)
    else:
        info("Launching new HashiCorp Vault container")
        run([
            "docker", "run", "-d",
            "--name", container_name,
            "-e", "VAULT_DISABLE_MLOCK=true",
            "--cap-add", "IPC_LOCK",
            "-p", "8200:8200",
            "-p", "8201:8201",
            "-p", "5696:5696",
            "-v", f"{certs_dir}:/vault/certs",
            "-v", f"{vault_hcl}:/vault/config/vault.hcl",
            "-v", f"{data_dir}:/vault/data",
            "-v", f"{log_dir}:/vault/log",
            "-e", f"VAULT_LICENSE={license_content}",
            "hashicorp/vault-enterprise:latest",
            "vault", "server", "-config=/vault/config/vault.hcl",
        ], capture_output=True)

    info("Waiting for Vault to start...")
    time.sleep(10)
    info("Vault container started")


def initialize_vault(container_name: str, certs_dir: Path) -> None:
    init_file = certs_dir / "init_token.json"

    if init_file.is_file():
        info(f"Vault already initialized at {init_file}")
    else:
        info("Initializing Vault...")

        # Capture init output on the host — avoids container write permission issues
        result = run(
            [
                "docker", "exec", container_name,
                "sh", "-c",
                'export VAULT_ADDR="http://127.0.0.1:8200" && vault operator init -format=json',
            ],
            capture_output=True,
        )
        init_file.write_text(result.stdout)
        info("Vault initialized.")

    info("Unsealing Vault...")
    init_data = json.loads(init_file.read_text())

    for i in range(3):
        unseal_key = init_data["unseal_keys_b64"][i]
        docker_exec(
            container_name,
            f"export VAULT_ADDR='http://127.0.0.1:8200' && vault operator unseal '{unseal_key}'",
            capture_output=True,
        )

    info("Vault unsealed successfully")


def configure_kmip(container_name: str, certs_dir: Path) -> None:
    init_file = certs_dir / "init_token.json"
    root_token = json.loads(init_file.read_text())["root_token"]

    info("Configuring KMIP secrets engine")
    vault_env = f"export VAULT_ADDR='http://127.0.0.1:8200' VAULT_TOKEN='{root_token}'"

    # Enable KMIP secrets engine and configure
    docker_exec(container_name, f"""
        {vault_env}
        set -e
        if ! vault secrets list | grep -q kmip/; then
            vault secrets enable kmip
        fi
        vault write kmip/config \\
            tls_ca_key_type='rsa' \\
            tls_ca_key_bits=2048 \\
            listen_addrs='0.0.0.0:5696' \\
            server_hostnames='172.17.0.1'
    """)

    # Read CA cert — capture on host to avoid container write permission issues
    result = docker_exec(
        container_name,
        f"{vault_env} && vault read -field=ca_pem kmip/ca",
        capture_output=True,
    )
    (certs_dir / "root_certificate.pem").write_text(result.stdout)

    # Create scope and role
    docker_exec(container_name, f"""
        {vault_env}
        set -e
        if ! vault list kmip/scope 2>/dev/null | grep -q my-service 2>/dev/null; then
            vault write -f kmip/scope/my-service
        fi
        vault write kmip/scope/my-service/role/admin \\
            operation_all=true \\
            tls_client_key_bits=2048 \\
            tls_client_key_type=rsa \\
            tls_client_ttl=24h
    """)

    # Generate credentials — capture on host to avoid container write permission issues
    result = docker_exec(
        container_name,
        f"{vault_env} && vault write -format=json kmip/scope/my-service/role/admin/credential/generate format=pem",
        capture_output=True,
    )
    credential_file = certs_dir / "credential.json"
    credential_file.write_text(result.stdout)

    credential_data = json.loads(result.stdout)
    (certs_dir / "client_certificate.pem").write_text(credential_data["data"]["certificate"])
    (certs_dir / "client_key.pem").write_text(credential_data["data"]["private_key"])

    info("KMIP configured successfully")


def verify_kmip_connection(certs_dir: Path) -> None:
    info("Verifying HashiCorp KMIP connection...")

    if not shutil.which("openssl"):
        error("openssl not found on host — skipping KMIP connection verification")
        return

    result = subprocess.run(
        [
            "openssl", "s_client",
            "-connect", "127.0.0.1:5696",
            "-CAfile", str(certs_dir / "root_certificate.pem"),
            "-cert",   str(certs_dir / "client_certificate.pem"),
            "-key",    str(certs_dir / "client_key.pem"),
            "-showcerts",
            "-status",
        ],
        input="",
        capture_output=True,
        text=True,
        timeout=10,
        check=False,
    )

    output = result.stdout + result.stderr
    if "Verification: OK" not in output and "Verify return code: 0" not in output:
        error("KMIP connection verification failed")
        sys.exit(1)

    info("KMIP connection verified successfully")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        prog="hashicorp-kmip-setup.py",
        description="Setup HashiCorp Vault Enterprise with KMIP support.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
LICENSE:
  Provide via environment variable (recommended):
    export VAULT_LICENSE=$(cat /path/to/vault.hclic)
    python3 hashicorp-kmip-setup.py

  Or via file path:
    python3 hashicorp-kmip-setup.py --license=/path/to/vault.hclic

EXAMPLES:
  python3 hashicorp-kmip-setup.py --license=/path/to/vault.hclic
  python3 hashicorp-kmip-setup.py --license=/path/to/vault.hclic --cert-dir=/custom/path/certs
  python3 hashicorp-kmip-setup.py --license=/path/to/vault.hclic --verbose
  export VAULT_LICENSE=$(cat vault.hclic) && python3 hashicorp-kmip-setup.py --verbose
        """,
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Enable verbose output",
    )
    parser.add_argument(
        "--cert-dir", "--certs-dir",
        dest="cert_dir",
        metavar="DIR",
        help="Directory where certificates will be stored (default: ~/vault/certs)",
    )
    parser.add_argument(
        "--license",
        required=False,
        metavar="FILE",
        help="Path to vault.hclic. Can also be set via: export VAULT_LICENSE=<content>",
    )
    return parser.parse_args()


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main() -> None:
    args = parse_args()

    # Resolve directories
    home = Path.home()
    if args.cert_dir and Path(args.cert_dir).is_dir():
        vault_base = Path(args.cert_dir) / "vault"
        certs_dir = Path(args.cert_dir)
    else:
        vault_base = home / "vault"
        certs_dir = vault_base / "certs"

    config_dir = vault_base / "config"
    data_dir   = vault_base / "data"
    log_dir    = vault_base / "log"
    vault_hcl  = config_dir / "vault.hcl"
    container_name = "kmip_hashicorp"

    setup_directories(config_dir, data_dir, log_dir, certs_dir)
    license_content = check_license(args.license)
    create_vault_hcl(vault_hcl)
    start_vault_container(container_name, certs_dir, vault_hcl, data_dir, log_dir, license_content)
    initialize_vault(container_name, certs_dir)
    configure_kmip(container_name, certs_dir)
    verify_kmip_connection(certs_dir)

    info("HashiCorp Vault Enterprise deployment successful")

    if args.verbose:
        info("HashiCorp KMIP setup completed successfully!")
        info(f"Files created within {certs_dir}:")
        info(f"  - Private key:      {certs_dir}/client_key.pem")
        info(f"  - Certificate:      {certs_dir}/client_certificate.pem")
        info(f"  - CA certificate:   {certs_dir}/root_certificate.pem")


if __name__ == "__main__":
    main()