#!/usr/bin/env python3
"""
==============================================================================
Elastic Stack Multi-Server Certificate Renewal & SCP Sync (Python)
==============================================================================
Handles dual-server topology:
  - Local Server: Runs Elasticsearch & Certificate Authority generation.
  - Remote Server: Runs Kibana and Elastic Agent (Fleet Server).
Transfers generated certificates to remote Kibana host via SCP, updates remote
kibana.yml fingerprint, restarts remote services, and re-enrolls remote Agent.
==============================================================================
"""

import os
import sys
import time
import shutil
import zipfile
import re
import subprocess
import urllib.request
import urllib.parse
import ssl
import json
import argparse

# Terminal ANSI Color & Formatting Definitions
RED = "\033[0;31m"
BOLD_RED = "\033[1;31m"
GREEN = "\033[0;32m"
BOLD_GREEN = "\033[1;32m"
YELLOW = "\033[1;33m"
BLUE = "\033[0;34m"
BOLD_BLUE = "\033[1;34m"
CYAN = "\033[0;36m"
BOLD_CYAN = "\033[1;36m"
DIM = "\033[2m"
RESET = "\033[0m"


def log_section(title):
    print(f"\n{BOLD_CYAN}===================================================={RESET}")
    print(f" {BOLD_BLUE}{title}{RESET}")
    print(f"{BOLD_CYAN}===================================================={RESET}")


def log_step(msg):
    print(f" {BOLD_CYAN}[+]{RESET} {msg}")


def log_sub(msg):
    print(f"    {CYAN}↳{RESET} {msg}")


def log_success(msg):
    print(f"    {BOLD_GREEN}✔{RESET} {msg}")


def log_warn(msg):
    print(f"    {YELLOW}⚠{RESET} {msg}")


def log_error(msg):
    print(f"    {BOLD_RED}✘{RESET} {msg}")


def spin_until(msg, check_fn, max_attempts=30, delay=2):
    spin_chars = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
    spin_i = 0
    for attempt in range(1, max_attempts + 1):
        if check_fn():
            print(f"\r    {BOLD_GREEN}✔{RESET} {msg} {BOLD_GREEN}(Ready!){RESET}                                  ")
            return True
        for _ in range(int(delay * 10)):
            char = spin_chars[spin_i]
            spin_i = (spin_i + 1) % len(spin_chars)
            print(f"\r    {YELLOW}{char}{RESET} {msg} {DIM}(attempt {attempt}/{max_attempts}){RESET}   ", end="", flush=True)
            time.sleep(0.1)
    print(f"\r    {BOLD_RED}✘{RESET} {msg} {BOLD_RED}(Timed out after {max_attempts}s){RESET}                                  ")
    return False


def run_cmd(cmd, check=False, capture=True, text=True):
    try:
        res = subprocess.run(cmd, check=check, stdout=subprocess.PIPE if capture else None, stderr=subprocess.PIPE if capture else None, text=text)
        return res
    except Exception as e:
        return subprocess.CompletedProcess(cmd, 1, "", str(e))


def write_instances_yaml(target_path, es_dns, es_ips, kib_dns, kib_ips):
    es_dns_str = "\n".join([f'      - "{d}"' for d in es_dns])
    es_ip_str = "\n".join([f'      - "{i}"' for i in es_ips])
    kib_dns_str = "\n".join([f'      - "{d}"' for d in kib_dns])
    kib_ip_str = "\n".join([f'      - "{i}"' for i in kib_ips])

    content = f"""instances:
  - name: "elasticsearch"
    dns:
{es_dns_str}
    ip:
{es_ip_str}
  - name: "kibana"
    dns:
{kib_dns_str}
    ip:
{kib_ip_str}
"""
    with open(target_path, "w", encoding="utf-8") as f:
        f.write(content)


def symlink_cert(src, dest):
    if src != dest:
        os.makedirs(os.path.dirname(dest), exist_ok=True)
        if os.path.exists(dest) or os.path.islink(dest):
            os.remove(dest)
        os.symlink(src, dest)
        log_sub(f"{src} ➜ {dest}")


def build_ssh_prefix(user, host, port=22, key_path=None):
    cmd = ["ssh", "-o", "StrictHostKeyChecking=no", "-p", str(port)]
    if key_path:
        cmd.extend(["-i", key_path])
    target = f"{user}@{host}" if user else host
    cmd.append(target)
    return cmd


def build_scp_prefix(port=22, key_path=None):
    cmd = ["scp", "-o", "StrictHostKeyChecking=no", "-P", str(port)]
    if key_path:
        cmd.extend(["-i", key_path])
    return cmd


def main():
    parser = argparse.ArgumentParser(description="Multi-Server Certificate Renewal & SCP Sync Script")
    parser.add_argument("--kibana-host", default=os.getenv("KIBANA_HOST", "10.1.10.4"), help="Remote Kibana host IP or hostname")
    parser.add_argument("--kibana-user", default=os.getenv("KIBANA_SSH_USER", "root"), help="SSH user for remote Kibana host")
    parser.add_argument("--kibana-port", default=os.getenv("KIBANA_SSH_PORT", "22"), help="SSH port for remote Kibana host")
    parser.add_argument("--ssh-key", default=os.getenv("KIBANA_SSH_KEY", ""), help="SSH private key path")
    args = parser.parse_args()

    if os.geteuid() != 0:
        print(f"{BOLD_RED}ERROR: Please run this script as root (using sudo){RESET}", file=sys.stderr)
        sys.exit(1)

    # Configuration values
    cert_dir = os.getenv("CERT_DIR", "/etc/elasticsearch/certs")
    es_http_p12_target = os.getenv("ES_HTTP_P12_TARGET", "/etc/elasticsearch/certs/http.p12")
    es_transport_p12_target = os.getenv("ES_TRANSPORT_P12_TARGET", "/etc/elasticsearch/certs/transport.p12")
    es_http_ca_target = os.getenv("ES_HTTP_CA_TARGET", "/etc/elasticsearch/certs/http_ca.crt")

    remote_kibana_host = args.kibana_host
    remote_kibana_user = args.kibana_user
    remote_kibana_port = args.kibana_port
    ssh_key = args.ssh_key

    remote_kibana_crt_target = os.getenv("REMOTE_KIBANA_CRT_TARGET", "/etc/kibana/certs/kibana.crt")
    remote_kibana_key_target = os.getenv("REMOTE_KIBANA_KEY_TARGET", "/etc/kibana/certs/kibana.key")
    remote_kibana_ca_target = os.getenv("REMOTE_KIBANA_CA_TARGET", "/var/lib/kibana/ca_1721297493415.crt")
    remote_kibana_config = os.getenv("REMOTE_KIBANA_CONFIG", "/etc/kibana/kibana.yml")

    fleet_server_es_url = os.getenv("FLEET_SERVER_ES_URL", "https://10.1.10.3:9200")
    fleet_server_url = os.getenv("FLEET_SERVER_URL", f"https://{remote_kibana_host}:8220")
    fleet_server_policy = os.getenv("FLEET_SERVER_POLICY", "fleet-server-policy")
    fleet_server_port = os.getenv("FLEET_SERVER_PORT", "8220")

    es_san_dns = [d for d in os.getenv("ES_SAN_DNS", "localhost cmsrv2024 siem2024").strip().split() if d]
    es_san_ips = [i for i in os.getenv("ES_SAN_IPS", "127.0.0.1 10.1.10.3").strip().split() if i]
    kib_san_dns = [d for d in os.getenv("KIBANA_SAN_DNS", f"localhost kibana-server {remote_kibana_host}").strip().split() if d]
    kib_san_ips = [i for i in os.getenv("KIBANA_SAN_IPS", f"127.0.0.1 {remote_kibana_host}").strip().split() if i]

    ca_path = os.path.join(cert_dir, "elastic-stack-ca.p12")
    es_certutil = "/usr/share/elasticsearch/bin/elasticsearch-certutil"

    ssh_base = build_ssh_prefix(remote_kibana_user, remote_kibana_host, remote_kibana_port, ssh_key)
    scp_base = build_scp_prefix(remote_kibana_port, ssh_key)
    remote_target_user = f"{remote_kibana_user}@{remote_kibana_host}" if remote_kibana_user else remote_kibana_host

    log_section(f"Multi-Server Elastic Stack Cert Renewal (Local: ES | Remote: {remote_kibana_host})")

    # Step 1: Local Certificate Generation
    log_step("Ensuring local certificate directory exists...")
    os.makedirs(cert_dir, exist_ok=True)
    log_sub(f"Local target directory: {cert_dir}")

    log_step("Checking Local Certificate Authority (CA)...")
    ca_valid = False
    if os.path.exists(ca_path):
        res = run_cmd(["openssl", "pkcs12", "-in", ca_path, "-passin", "pass:", "-nokeys"])
        if res.returncode == 0:
            log_success("Existing passwordless CA found.")
            ca_valid = True
        else:
            log_warn("CA exists but is password-protected or invalid.")

    if not ca_valid:
        log_step("Regenerating fresh passwordless CA...")
        if os.path.exists(ca_path):
            os.remove(ca_path)
        run_cmd([es_certutil, "ca", "--out", ca_path, "--pass", ""])
        log_success("New CA generated successfully.")

    log_step("Extracting HTTP CA certificate...")
    http_ca_temp = os.path.join(cert_dir, "http_ca.crt")
    run_cmd(["openssl", "pkcs12", "-in", ca_path, "-passin", "pass:", "-nokeys", "-out", http_ca_temp])
    log_success(f"Extracted HTTP CA certificate to {http_ca_temp}")

    log_step("Generating multi-server instances configuration...")
    instances_yml_path = os.path.join(cert_dir, "instances.yml")
    write_instances_yaml(instances_yml_path, es_san_dns, es_san_ips, kib_san_dns, kib_san_ips)
    log_sub("instances.yml written.")

    log_step("Generating node certificates for Elasticsearch and Kibana...")
    certs_zip_path = os.path.join(cert_dir, "certs.zip")
    if os.path.exists(certs_zip_path):
        os.remove(certs_zip_path)
    run_cmd([es_certutil, "cert", "--ca", ca_path, "--ca-pass", "", "--in", instances_yml_path, "--out", certs_zip_path, "--pass", ""])
    log_success("Node certificates generated into certs.zip")

    log_step("Extracting certificate archive...")
    if os.path.exists(certs_zip_path):
        with zipfile.ZipFile(certs_zip_path, "r") as zip_ref:
            zip_ref.extractall(cert_dir)
        log_sub("Certificates extracted.")

    log_step("Processing local Elasticsearch certificates...")
    es_p12_extracted = os.path.join(cert_dir, "elasticsearch", "elasticsearch.p12")
    http_p12_path = os.path.join(cert_dir, "http.p12")
    transport_p12_path = os.path.join(cert_dir, "transport.p12")

    if os.path.exists(es_p12_extracted):
        shutil.move(es_p12_extracted, http_p12_path)
        shutil.copy(http_p12_path, transport_p12_path)

    kibana_p12_extracted = os.path.join(cert_dir, "kibana", "kibana.p12")
    kibana_crt_path = os.path.join(cert_dir, "kibana.crt")
    kibana_key_tmp = os.path.join(cert_dir, "kibana.key.tmp")
    kibana_key_path = os.path.join(cert_dir, "kibana.key")

    if os.path.exists(kibana_p12_extracted):
        run_cmd(["openssl", "pkcs12", "-in", kibana_p12_extracted, "-clcerts", "-nokeys", "-out", kibana_crt_path, "-passin", "pass:"])
        run_cmd(["openssl", "pkcs12", "-in", kibana_p12_extracted, "-nocerts", "-nodes", "-out", kibana_key_tmp, "-passin", "pass:"])
        run_cmd(["openssl", "pkey", "-in", kibana_key_tmp, "-out", kibana_key_path])

    for cleanup_item in [os.path.join(cert_dir, "elasticsearch"), os.path.join(cert_dir, "kibana"), certs_zip_path, instances_yml_path, kibana_key_tmp]:
        if os.path.exists(cleanup_item):
            if os.path.isdir(cleanup_item):
                shutil.rmtree(cleanup_item)
            else:
                os.remove(cleanup_item)

    log_step("Clearing legacy secure keystore passwords on local Elasticsearch...")
    es_keystore = "/usr/share/elasticsearch/bin/elasticsearch-keystore"
    for key in ["xpack.security.http.ssl.keystore.secure_password", "xpack.security.transport.ssl.keystore.secure_password", "xpack.security.transport.ssl.truststore.secure_password", "xpack.security.http.ssl.truststore.secure_password"]:
        run_cmd([es_keystore, "remove", key])

    log_step("Applying permissions & updating local Elasticsearch symlinks...")
    symlink_cert(http_p12_path, es_http_p12_target)
    symlink_cert(transport_p12_path, es_transport_p12_target)
    symlink_cert(http_ca_temp, es_http_ca_target)

    run_cmd(["systemctl", "restart", "elasticsearch"])
    spin_until("Waiting for local Elasticsearch on port 9200...", lambda: run_cmd(["curl", "-k", "-s", "https://localhost:9200"]).returncode == 0, max_attempts=30, delay=6)
    log_success("Local Elasticsearch restarted and healthy.")

    # Step 2: SCP Certificates to Remote Kibana Server
    log_section(f"Transferring Certificates via SCP to Remote Host ({remote_kibana_host})")

    log_step(f"Ensuring remote directories exist on {remote_kibana_host}...")
    remote_crt_dir = os.path.dirname(remote_kibana_crt_target)
    remote_ca_dir = os.path.dirname(remote_kibana_ca_target)

    run_cmd(ssh_base + [f"mkdir -p {remote_crt_dir} {remote_ca_dir} /tmp/elastic-certs-staging"])
    log_sub("Remote directories prepared.")

    log_step("SCP Kibana certificate, key, and CA to remote host...")
    run_cmd(scp_base + [kibana_crt_path, f"{remote_target_user}:/tmp/elastic-certs-staging/kibana.crt"])
    run_cmd(scp_base + [kibana_key_path, f"{remote_target_user}:/tmp/elastic-certs-staging/kibana.key"])
    run_cmd(scp_base + [http_ca_temp, f"{remote_target_user}:/tmp/elastic-certs-staging/http_ca.crt"])
    log_success("Certificates uploaded to remote staging area.")

    log_step("Installing certificates into target remote locations & setting permissions...")
    remote_install_cmd = (
        f"cp /tmp/elastic-certs-staging/kibana.crt {remote_kibana_crt_target} && "
        f"cp /tmp/elastic-certs-staging/kibana.key {remote_kibana_key_target} && "
        f"cp /tmp/elastic-certs-staging/http_ca.crt {remote_kibana_ca_target} && "
        f"chown -R root:kibana {remote_crt_dir} && "
        f"chmod 644 {remote_kibana_crt_target} && "
        f"chmod 640 {remote_kibana_key_target} && "
        f"chmod 644 {remote_kibana_ca_target} && "
        f"rm -rf /tmp/elastic-certs-staging"
    )
    run_cmd(ssh_base + [remote_install_cmd])
    log_success("Certificates installed on remote Kibana host.")

    # Step 3: Compute CA Fingerprint & Update Remote Kibana Config
    log_step("Computing CA fingerprint and updating remote Kibana config...")
    new_fingerprint = ""
    res = run_cmd(["openssl", "x509", "-in", http_ca_temp, "-noout", "-sha256", "-fingerprint"])
    if res.returncode == 0:
        match = re.search(r"Fingerprint=(.*)", res.stdout)
        if match:
            new_fingerprint = match.group(1).replace(":", "").lower().strip()

    log_sub(f"New CA Fingerprint: {BOLD_GREEN}{new_fingerprint}{RESET}")

    if new_fingerprint:
        update_remote_kibana = (
            f"python3 -c \"import re; f=open('{remote_kibana_config}', 'r').read(); "
            f"f=re.sub(r'ca_trusted_fingerprint:\\s*[a-fA-F0-9]+', 'ca_trusted_fingerprint: {new_fingerprint}', f); "
            f"open('{remote_kibana_config}', 'w').write(f)\" || "
            f"sed -i 's/ca_trusted_fingerprint:.*/ca_trusted_fingerprint: {new_fingerprint}/' {remote_kibana_config}"
        )
        run_cmd(ssh_base + [update_remote_kibana])
        log_success("Remote Kibana config updated with new CA fingerprint.")

    log_step("Restarting remote Kibana service...")
    run_cmd(ssh_base + ["systemctl restart kibana"])

    def check_remote_kibana():
        res = run_cmd(ssh_base + ["curl -k -s https://localhost:5601"])
        return res.returncode == 0

    spin_until(f"Waiting for remote Kibana on {remote_kibana_host}:5601...", check_remote_kibana)

    # Step 4: Re-enroll Remote Elastic Agent (Fleet Server)
    log_section("Re-enrolling Remote Elastic Agent (Fleet Server)")

    log_step("Generating Fleet Server service token from local Elasticsearch...")
    token_name = f"cert-exchange-remote-{int(time.time())}"
    fleet_service_token = ""

    for attempt in range(1, 13):
        res = run_cmd(["/usr/share/elasticsearch/bin/elasticsearch-service-tokens", "create", "elastic/fleet-server", token_name])
        if res.returncode == 0 and res.stdout.strip():
            fleet_service_token = res.stdout.strip().split()[-1]
            break
        time.sleep(5)

    if not fleet_service_token:
        log_error("Failed to generate Fleet Server service token.")
        sys.exit(1)
    log_success("Generated Fleet Server service token.")

    log_step("Stopping remote Elastic Agent service...")
    run_cmd(ssh_base + ["systemctl stop elastic-agent || true"])

    log_step("Re-enrolling remote Elastic Agent (Fleet Server)...")
    remote_enroll_cmd = (
        f"/opt/Elastic/Agent/elastic-agent enroll --force "
        f"--url={fleet_server_url} "
        f"--fleet-server-es={fleet_server_es_url} "
        f"--fleet-server-es-ca={remote_kibana_ca_target} "
        f"--fleet-server-service-token={fleet_service_token} "
        f"--fleet-server-policy={fleet_server_policy} "
        f"--fleet-server-es-ca-trusted-fingerprint={new_fingerprint} "
        f"--fleet-server-port={fleet_server_port} "
        f"--insecure"
    )
    run_cmd(ssh_base + [remote_enroll_cmd])
    log_success("Remote Elastic Agent re-enrolled.")

    log_step("Starting remote Elastic Agent service...")
    run_cmd(ssh_base + ["systemctl start elastic-agent"])

    def check_remote_fleet():
        res = run_cmd(ssh_base + ["curl -k -s https://localhost:8220"])
        return res.returncode == 0

    spin_until(f"Waiting for remote Fleet Server on {remote_kibana_host}:8220...", check_remote_fleet)

    log_section("Multi-Server SSL Verification")
    log_step(f"Remote Kibana ({remote_kibana_host}:5601):")
    r1 = run_cmd(["openssl", "s_client", "-connect", f"{remote_kibana_host}:5601", "-showcerts"], text=True)
    if r1.returncode == 0 and "END CERTIFICATE" in r1.stdout:
        log_success("Remote Kibana SSL endpoint verified.")
    else:
        log_warn("Could not verify remote Kibana endpoint directly (check firewall).")

    log_step(f"Remote Fleet Server ({remote_kibana_host}:8220):")
    r2 = run_cmd(["openssl", "s_client", "-connect", f"{remote_kibana_host}:8220", "-showcerts"], text=True)
    if r2.returncode == 0 and "END CERTIFICATE" in r2.stdout:
        log_success("Remote Fleet Server SSL endpoint verified.")
    else:
        log_warn("Could not verify remote Fleet Server endpoint directly (check firewall).")

    print(f"\n{BOLD_GREEN}===================================================={RESET}")
    print(f" {BOLD_GREEN}✔ Multi-server certificate exchange completed!{RESET}")
    print(f"{BOLD_GREEN}===================================================={RESET}\n")


if __name__ == "__main__":
    main()