#!/usr/bin/env python3
"""
fix_cdw_ssl.py — Fix SSL PKIX error in the databus-producer Kubernetes deployment.

Downloads the server certificate from the telemetry.altus.url configured in the
databus-producer configmap, builds a JKS truststore seeded with the pod's own system
CA bundle (so all existing HTTPS traffic stays working), creates a Kubernetes secret,
and patches the deployment to mount the truststore and set JVM_OPTS.

Usage:
    export KUBECONFIG=/path/to/kubeconfig.yaml
    python3 fix_cdw_ssl.py --namespace <namespace>

Requirements (runtime): Python 3.8+, kubectl on PATH, keytool (JDK) on PATH
"""

import argparse
import json
import os
import shutil
import socket
import ssl
import subprocess
import sys
import tempfile
import xml.etree.ElementTree as ET
from datetime import datetime
from typing import Callable, Dict, List, Optional
from urllib.parse import urlparse

# ---------------------------------------------------------------------------
# Exception hierarchy
# ---------------------------------------------------------------------------

class DatabusSSLError(Exception):
    """Base error for all script failures."""


class PrerequisiteError(DatabusSSLError):
    """Missing tool on PATH, KUBECONFIG not set or file missing."""


class DeploymentError(DatabusSSLError):
    """Deployment not found, already patched, pod not running, or secret exists."""


class ConfigMapError(DatabusSSLError):
    """ConfigMap missing, telemetry.altus.url property missing or empty."""


class CertificateError(DatabusSSLError):
    """Certificate download failure or no CA bundle found in pod."""


class TruststoreError(DatabusSSLError):
    """keytool failure while building the truststore."""


class BackupError(DatabusSSLError):
    """Backup export failed — no K8s resources were mutated."""


class RollbackError(DatabusSSLError):
    """Rollback failed (logged only, not re-raised)."""


# ---------------------------------------------------------------------------
# Pure functions — no I/O, fully unit-testable without mocks
# ---------------------------------------------------------------------------

def parse_telemetry_url(xml_content: str) -> str:
    """Parse telemetry.altus.url from databus-producer.xml content.

    Args:
        xml_content: Raw XML string from the configmap.

    Returns:
        The URL string (e.g. 'https://host:12022').

    Raises:
        ConfigMapError: If the property is missing or its value is empty.
    """
    root = ET.fromstring(xml_content)
    for prop in root.findall("property"):
        name_el = prop.find("name")
        if name_el is not None and name_el.text == "telemetry.altus.url":
            value_el = prop.find("value")
            if value_el is None or not (value_el.text or "").strip():
                raise ConfigMapError(
                    "telemetry.altus.url property found but value is empty"
                )
            return value_el.text.strip()
    raise ConfigMapError(
        "telemetry.altus.url property not found in databus-producer.xml"
    )


def parse_pem_certificates(pem_bundle: str) -> List[str]:
    """Split a PEM bundle string into individual certificate strings.

    Args:
        pem_bundle: String containing one or more PEM certificate blocks.

    Returns:
        List of individual PEM certificate strings (each includes header/footer).
        Returns empty list if pem_bundle contains no certificate blocks.
    """
    certs = []
    current: List[str] = []
    in_cert = False
    for line in pem_bundle.splitlines():
        if line.strip() == "-----BEGIN CERTIFICATE-----":
            in_cert = True
            current = [line]
        elif line.strip() == "-----END CERTIFICATE-----":
            current.append(line)
            certs.append("\n".join(current) + "\n")
            current = []
            in_cert = False
        elif in_cert:
            current.append(line)
    return certs


# ---------------------------------------------------------------------------
# I/O boundary functions — all external calls injectable via `run` parameter
# ---------------------------------------------------------------------------

def validate_prerequisites(env: Dict[str, str], run: Callable) -> None:
    """Validate all prerequisites before any K8s or network calls.

    Checks:
    - KUBECONFIG env var is set
    - KUBECONFIG file exists on disk
    - kubectl is on PATH
    - keytool is on PATH

    Args:
        env: Environment variable dict (pass os.environ in production).
        run: Callable matching subprocess.run signature (injectable for tests).

    Raises:
        PrerequisiteError: If any prerequisite is not met.
    """
    kubeconfig = env.get("KUBECONFIG")
    if not kubeconfig:
        raise PrerequisiteError(
            "KUBECONFIG environment variable is not set. "
            "Run: export KUBECONFIG=/path/to/kubeconfig.yaml"
        )
    if not os.path.isfile(kubeconfig):
        raise PrerequisiteError(
            f"KUBECONFIG file not found: {kubeconfig}"
        )
    if not shutil.which("kubectl"):
        raise PrerequisiteError(
            "kubectl not found on PATH. Install kubectl and ensure it is on your PATH."
        )
    if not shutil.which("keytool"):
        raise PrerequisiteError(
            "keytool not found on PATH. Install a JDK (not JRE) and ensure it is on your PATH."
        )


def find_deployment(namespace: str, run: Callable) -> dict:
    """Retrieve the databus-producer Deployment as a parsed JSON dict.

    Args:
        namespace: Kubernetes namespace.
        run: subprocess.run-compatible callable.

    Returns:
        Parsed deployment JSON dict.

    Raises:
        DeploymentError: If the deployment is not found.
    """
    result = run(
        ["kubectl", "get", "deployment", "databus-producer",
         "-n", namespace, "-o", "json"],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise DeploymentError(
            f"databus-producer deployment not found in namespace '{namespace}'. "
            f"Error: {result.stderr.strip()}"
        )
    return json.loads(result.stdout)


def check_jvm_opts_not_set(deployment: dict) -> None:
    """Verify JVM_OPTS is not already set in the databus-producer container.

    Args:
        deployment: Parsed deployment JSON dict (from find_deployment).

    Raises:
        DeploymentError: If JVM_OPTS is already set.
    """
    containers = (
        deployment.get("spec", {})
        .get("template", {})
        .get("spec", {})
        .get("containers", [])
    )
    for container in containers:
        if container.get("name") == "databus-producer":
            for env_var in container.get("env", []):
                if env_var.get("name") == "JVM_OPTS":
                    raise DeploymentError(
                        "JVM_OPTS is already set in the databus-producer container. "
                        "This script should not be run on an already-patched deployment."
                    )


def check_secret_not_exists(namespace: str, run: Callable) -> None:
    """Verify the databus-producer-truststore secret does not already exist.

    Args:
        namespace: Kubernetes namespace.
        run: subprocess.run-compatible callable.

    Raises:
        DeploymentError: If the secret already exists.
    """
    result = run(
        ["kubectl", "get", "secret", "databus-producer-truststore",
         "-n", namespace],
        capture_output=True, text=True,
    )
    if result.returncode == 0:
        raise DeploymentError(
            "Secret 'databus-producer-truststore' already exists in namespace "
            f"'{namespace}'. This script should not be run again on an already-patched deployment."
        )


def backup_resources(namespace: str, backup_dir: str, run: Callable) -> None:
    """Export deployment and configmap YAMLs to a local backup directory.

    Writes:
    - <backup_dir>/deployment-databus-producer.yaml
    - <backup_dir>/configmap-databus-producer-conf.yaml

    Args:
        namespace: Kubernetes namespace.
        backup_dir: Path to (already created) backup directory.
        run: subprocess.run-compatible callable.

    Raises:
        BackupError: If any export fails.
    """
    resources = [
        (
            ["kubectl", "get", "deployment", "databus-producer",
             "-n", namespace, "-o", "yaml"],
            "deployment-databus-producer.yaml",
        ),
        (
            ["kubectl", "get", "configmap", "databus-producer-conf-databus-producer",
             "-n", namespace, "-o", "yaml"],
            "configmap-databus-producer-conf.yaml",
        ),
    ]
    for cmd, filename in resources:
        result = run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            raise BackupError(
                f"Failed to back up {filename}: {result.stderr.strip()}. "
                "No Kubernetes resources have been modified."
            )
        out_path = os.path.join(backup_dir, filename)
        with open(out_path, "w") as f:
            f.write(result.stdout)


def get_configmap(namespace: str, run: Callable) -> str:
    """Retrieve the databus-producer.xml content from the configmap.

    Args:
        namespace: Kubernetes namespace.
        run: subprocess.run-compatible callable.

    Returns:
        Raw XML string from the databus-producer.xml key.

    Raises:
        ConfigMapError: If the configmap or the key is not found.
    """
    result = run(
        ["kubectl", "get", "configmap", "databus-producer-conf-databus-producer",
         "-n", namespace, "-o", "json"],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise ConfigMapError(
            f"ConfigMap 'databus-producer-conf-databus-producer' not found in "
            f"namespace '{namespace}'. Error: {result.stderr.strip()}"
        )
    cm = json.loads(result.stdout)
    xml_content = cm.get("data", {}).get("databus-producer.xml")
    if xml_content is None:
        raise ConfigMapError(
            "Key 'databus-producer.xml' not found in configmap "
            "'databus-producer-conf-databus-producer'."
        )
    return xml_content


def check_pod_running(namespace: str, run: Callable) -> str:
    """Find a Running databus-producer pod and return its name.

    Args:
        namespace: Kubernetes namespace.
        run: subprocess.run-compatible callable.

    Returns:
        Pod name string.

    Raises:
        DeploymentError: If no databus-producer pod is in Running state.
    """
    result = run(
        ["kubectl", "get", "pods", "-n", namespace,
         "-l", "app=databus-producer", "-o", "json"],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise DeploymentError(
            f"Failed to list pods in namespace '{namespace}': {result.stderr.strip()}"
        )
    pods_data = json.loads(result.stdout)
    for pod in pods_data.get("items", []):
        phase = pod.get("status", {}).get("phase", "")
        if phase == "Running":
            return pod["metadata"]["name"]
    raise DeploymentError(
        f"No databus-producer pod found in Running state in namespace '{namespace}'. "
        "Ensure the pod is running before executing this script."
    )


def _fetch_pem_certificate(host: str, port: int) -> str:
    """Open an SSL connection and return the server certificate as PEM string.

    Separated from download_server_certificate so tests can patch it directly.
    """
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    with ctx.wrap_socket(socket.create_connection((host, port), timeout=10),
                         server_hostname=host) as ssock:
        der_cert = ssock.getpeercert(binary_form=True)
    pem = ssl.DER_cert_to_PEM_cert(der_cert)
    return pem


def download_server_certificate(host: str, port: int, temp_dir: str) -> str:
    """Download the server's TLS certificate via an SSL socket connection.

    Args:
        host: Hostname to connect to.
        port: TCP port.
        temp_dir: Directory where the PEM file will be written.

    Returns:
        Absolute path to the written PEM file.

    Raises:
        CertificateError: If the connection fails or no certificate is returned.
    """
    try:
        pem = _fetch_pem_certificate(host, port)
    except (OSError, socket.error) as exc:
        raise CertificateError(
            f"Could not connect to {host}:{port} to download the server certificate. "
            f"Ensure you have network/VPN access to this host. Error: {exc}"
        ) from exc

    if not pem or "BEGIN CERTIFICATE" not in pem:
        raise CertificateError(
            f"No certificate returned from {host}:{port}."
        )

    cert_path = os.path.join(temp_dir, "server.pem")
    with open(cert_path, "w") as f:
        f.write(pem)
    return cert_path


_CA_BUNDLE_PATHS = [
    "/etc/ssl/certs/ca-certificates.crt",   # Debian/Ubuntu
    "/etc/pki/tls/certs/ca-bundle.crt",     # RHEL/CentOS
]


def get_system_ca_bundle_from_pod(pod_name: str, namespace: str, run: Callable) -> str:
    """Read the system CA bundle from inside the running databus-producer pod.

    Probes these paths in order:
    - /etc/ssl/certs/ca-certificates.crt  (Debian/Ubuntu)
    - /etc/pki/tls/certs/ca-bundle.crt    (RHEL/CentOS)

    Args:
        pod_name: Name of the running databus-producer pod.
        namespace: Kubernetes namespace.
        run: subprocess.run-compatible callable.

    Returns:
        PEM bundle string from the first path that succeeds.

    Raises:
        CertificateError: If neither path is found in the pod.
    """
    for path in _CA_BUNDLE_PATHS:
        result = run(
            ["kubectl", "exec", pod_name, "-n", namespace,
             "-c", "databus-producer", "--", "cat", path],
            capture_output=True, text=True,
        )
        if result.returncode == 0 and result.stdout.strip():
            return result.stdout

    raise CertificateError(
        f"No system CA bundle found in pod '{pod_name}'. "
        f"Checked paths: {', '.join(_CA_BUNDLE_PATHS)}. "
        "Cannot build a safe truststore without the pod's system CA certificates."
    )


def build_truststore(
    system_pems: List[str],
    server_cert_path: str,
    truststore_path: str,
    run: Callable,
) -> None:
    """Build a JKS truststore containing system CAs and the server certificate.

    Imports each system CA with alias 'system-ca-N' (index-based to avoid
    duplicate-subject collisions), then imports the server cert last with
    alias 'databus-producer-server'.

    Args:
        system_pems: List of individual PEM certificate strings from the pod.
        server_cert_path: Path to the downloaded server PEM file.
        truststore_path: Destination path for the truststore.jks file.
        run: subprocess.run-compatible callable.

    Raises:
        TruststoreError: If any keytool invocation fails.
    """
    temp_dir = os.path.dirname(truststore_path)

    def _keytool_import(cert_path: str, alias: str) -> None:
        cmd = [
            "keytool", "-importcert",
            "-noprompt",
            "-keystore", truststore_path,
            "-storepass", "changeit",
            "-alias", alias,
            "-file", cert_path,
        ]
        result = run(cmd, capture_output=True, text=True)
        if result.returncode != 0:
            raise TruststoreError(
                f"keytool failed to import certificate with alias '{alias}': "
                f"{result.stderr.strip()}"
            )

    # Import system CAs first
    for i, pem in enumerate(system_pems):
        pem_file = os.path.join(temp_dir, f"system-ca-{i}.pem")
        with open(pem_file, "w") as f:
            f.write(pem)
        _keytool_import(pem_file, f"system-ca-{i}")

    # Import server cert last
    _keytool_import(server_cert_path, "databus-producer-server")


def create_k8s_secret(namespace: str, truststore_path: str, run: Callable) -> None:
    """Create the databus-producer-truststore Kubernetes secret.

    Args:
        namespace: Kubernetes namespace.
        truststore_path: Path to the truststore.jks file.
        run: subprocess.run-compatible callable.

    Raises:
        DatabusSSLError: If kubectl create secret fails.
    """
    result = run(
        [
            "kubectl", "create", "secret", "generic",
            "databus-producer-truststore",
            f"--from-file=truststore.jks={truststore_path}",
            "-n", namespace,
        ],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise DatabusSSLError(
            f"Failed to create Kubernetes secret: {result.stderr.strip()}"
        )


_TRUSTSTORE_MOUNT_PATH = "/etc/databus-producer/trust"
_TRUSTSTORE_SECRET_NAME = "databus-producer-truststore"
_JVM_OPTS_VALUE = (
    "-Djavax.net.ssl.trustStore=/etc/databus-producer/trust/truststore.jks "
    "-Djavax.net.ssl.trustStorePassword=changeit"
)

_DEPLOYMENT_PATCH = {
    "spec": {
        "template": {
            "spec": {
                "volumes": [
                    {
                        "name": "truststore-volume",
                        "secret": {"secretName": _TRUSTSTORE_SECRET_NAME},
                    }
                ],
                "containers": [
                    {
                        "name": "databus-producer",
                        "env": [
                            {"name": "JVM_OPTS", "value": _JVM_OPTS_VALUE}
                        ],
                        "volumeMounts": [
                            {
                                "name": "truststore-volume",
                                "mountPath": _TRUSTSTORE_MOUNT_PATH,
                                "readOnly": True,
                            }
                        ],
                    }
                ],
            }
        }
    }
}


def patch_deployment(namespace: str, run: Callable) -> None:
    """Patch the databus-producer deployment to mount the truststore secret
    and set JVM_OPTS on the databus-producer container.

    Adds:
    - A volume referencing the databus-producer-truststore secret
    - A volumeMount at /etc/databus-producer/trust/ on the databus-producer container
    - JVM_OPTS env var on the databus-producer container only

    Args:
        namespace: Kubernetes namespace.
        run: subprocess.run-compatible callable.

    Raises:
        DeploymentError: If the kubectl patch fails.
    """
    result = run(
        [
            "kubectl", "patch", "deployment", "databus-producer",
            "-n", namespace,
            "--type", "strategic",
            "-p", json.dumps(_DEPLOYMENT_PATCH),
        ],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        raise DeploymentError(
            f"Failed to patch deployment: {result.stderr.strip()}"
        )


def rollback_secret(namespace: str, run: Callable) -> None:
    """Delete the databus-producer-truststore secret as a rollback step.

    Best-effort: logs a warning if deletion fails but does not re-raise,
    so the caller can still print the full rollback instructions to the user.

    Args:
        namespace: Kubernetes namespace.
        run: subprocess.run-compatible callable.
    """
    result = run(
        ["kubectl", "delete", "secret", "databus-producer-truststore",
         "-n", namespace],
        capture_output=True, text=True,
    )
    if result.returncode != 0:
        print(
            f"WARNING: Failed to delete secret 'databus-producer-truststore' "
            f"during rollback: {result.stderr.strip()}",
            file=sys.stderr,
        )


def print_success_summary(namespace: str, backup_dir: str) -> None:
    """Print the backup path and rollback instructions to stdout.

    Args:
        namespace: Kubernetes namespace (used in rollback commands).
        backup_dir: Path to the timestamped backup directory.
    """
    print("\n" + "=" * 60)
    print("SUCCESS: databus-producer SSL truststore configured.")
    print("=" * 60)
    print(f"\nBackup saved to: {backup_dir}")
    print("\nTo revert all changes:")
    print(
        f"  kubectl apply -f {backup_dir}/deployment-databus-producer.yaml "
        f"-n {namespace}"
    )
    print(
        f"  kubectl delete secret databus-producer-truststore -n {namespace}"
    )
    print()


# ---------------------------------------------------------------------------
# CLI entry point
# ---------------------------------------------------------------------------

def parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace:
    """Parse command-line arguments.

    Args:
        argv: Argument list (defaults to sys.argv[1:]).

    Returns:
        Parsed namespace with .namespace attribute.
    """
    parser = argparse.ArgumentParser(
        description="Fix SSL PKIX error in the databus-producer deployment."
    )
    parser.add_argument(
        "--namespace", "-n",
        required=True,
        help="Kubernetes namespace containing the databus-producer deployment.",
    )
    return parser.parse_args(argv)


def main() -> None:
    """Orchestrate all steps in order with full error handling and rollback."""
    args = parse_args()
    namespace = args.namespace

    # Validate tools and environment before touching anything
    validate_prerequisites(env=os.environ, run=subprocess.run)

    # Create timestamped backup directory (created before any mutations)
    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    backup_dir = os.path.abspath(f"databus-ssl-backup-{timestamp}")

    try:
        with tempfile.TemporaryDirectory() as temp_dir:
            # --- Guard checks (read-only) ---
            deployment = find_deployment(namespace, run=subprocess.run)
            check_jvm_opts_not_set(deployment)
            check_secret_not_exists(namespace, run=subprocess.run)

            # --- Read configmap and extract URL ---
            xml_content = get_configmap(namespace, run=subprocess.run)
            telemetry_url = parse_telemetry_url(xml_content)

            # Parse host and port from URL (e.g. https://host:12022)
            parsed = urlparse(telemetry_url)
            host = parsed.hostname
            port = parsed.port or 443

            # --- Backup before any mutations ---
            os.makedirs(backup_dir, exist_ok=True)
            backup_resources(namespace, backup_dir, run=subprocess.run)
            print(f"[OK] Backup saved to: {backup_dir}")

            # --- Pod must be running for kubectl exec ---
            pod_name = check_pod_running(namespace, run=subprocess.run)
            print(f"[OK] Found running pod: {pod_name}")

            # --- Download server certificate ---
            print(f"[..] Downloading server certificate from {host}:{port} ...")
            server_cert_path = download_server_certificate(host, port, temp_dir)
            print(f"[OK] Server certificate saved to: {server_cert_path}")

            # --- Read system CA bundle from pod ---
            print(f"[..] Reading system CA bundle from pod {pod_name} ...")
            ca_bundle = get_system_ca_bundle_from_pod(pod_name, namespace, run=subprocess.run)
            system_pems = parse_pem_certificates(ca_bundle)
            print(f"[OK] Found {len(system_pems)} system CA certificates in pod.")

            # --- Build truststore ---
            truststore_path = os.path.join(temp_dir, "truststore.jks")
            print(f"[..] Building truststore with {len(system_pems)} system CAs + server cert ...")
            build_truststore(system_pems, server_cert_path, truststore_path, run=subprocess.run)
            print(f"[OK] truststore.jks created.")

            # --- Create Kubernetes secret ---
            print("[..] Creating Kubernetes secret 'databus-producer-truststore' ...")
            create_k8s_secret(namespace, truststore_path, run=subprocess.run)
            print("[OK] Secret created.")

            # --- Patch deployment ---
            print("[..] Patching databus-producer deployment ...")
            try:
                patch_deployment(namespace, run=subprocess.run)
            except DeploymentError as exc:
                print(f"\nERROR: Deployment patch failed: {exc}", file=sys.stderr)
                print("Rolling back: deleting secret ...", file=sys.stderr)
                rollback_secret(namespace, run=subprocess.run)
                print(
                    f"\nTo restore the deployment, run:\n"
                    f"  kubectl apply -f {backup_dir}/deployment-databus-producer.yaml "
                    f"-n {namespace}",
                    file=sys.stderr,
                )
                raise
            print("[OK] Deployment patched.")

            print_success_summary(namespace, backup_dir)

    except DatabusSSLError as exc:
        print(f"\nERROR: {exc}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()