#!/usr/bin/env python3
"""
==============================================================================
Elastic Stack Certificate Exchange Autopilot (Self-Contained Single Script)
==============================================================================
Zero-interaction, full-stack certificate renewal & sync orchestrator.
Supports:
  - Single-Host Topology (ES + Kibana + Agent on 1 node)
  - Dual-Elasticsearch Nodes (--es2-host)
  - Remote Kibana + Fleet Server Agent (--kibana-host)

One-Liner Execution Examples:
  1. Single-Server:
     curl -sSL http://static.maixnor.com/certs/autopilot-certs.py | sudo python3 -

  2. Dual Elasticsearch Nodes:
     curl -sSL http://static.maixnor.com/certs/autopilot-certs.py | sudo python3 - --es2-host 10.1.10.2

  3. Dual Elasticsearch Nodes + Separate Kibana Server:
     curl -sSL http://static.maixnor.com/certs/autopilot-certs.py | sudo python3 - --es2-host 10.1.10.2 --kibana-host 10.1.10.4
==============================================================================
"""

import os
import sys
import time
import shutil
import zipfile
import re
import subprocess
import socket
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[0;33m"
BOLD_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"
BOLD = "\033[1m"
RESET = "\033[0m"


def log_header(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
    total_timeout = int(max_attempts * delay)
    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 {total_timeout}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 parse_simple_yaml(file_path):
    if not os.path.exists(file_path):
        return {}
    config = {}
    try:
        with open(file_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                if ":" in line:
                    parts = line.split(":", 1)
                    k = parts[0].strip()
                    v = parts[1].strip().strip("\"'")
                    if v:
                        config[k] = v
    except Exception:
        pass
    return config


def discover_network_sans():
    dns_names = set(["localhost"])
    ip_addresses = set(["127.0.0.1"])

    try:
        hostname = socket.gethostname()
        if hostname:
            dns_names.add(hostname)
            fqdn = socket.getfqdn()
            if fqdn:
                dns_names.add(fqdn)
    except Exception:
        pass

    try:
        res = subprocess.run(["hostname", "-I"], capture_output=True, text=True, timeout=2)
        if res.returncode == 0:
            for ip in res.stdout.strip().split():
                if ":" not in ip:
                    ip_addresses.add(ip)
    except Exception:
        pass

    for default_dns in ["cmsrv2024", "siem2024"]:
        dns_names.add(default_dns)
    for default_ip in ["10.1.10.3", "192.168.196.95"]:
        ip_addresses.add(default_ip)

    return sorted(list(dns_names)), sorted(list(ip_addresses))


def write_instances_yaml(target_path, es1_dns, es1_ips, es2_host=None, es2_dns=None, es2_ips=None, kib_dns=None, kib_ips=None):
    if kib_dns is None:
        kib_dns = es1_dns
    if kib_ips is None:
        kib_ips = es1_ips

    es1_dns_str = "\n".join([f'      - "{d}"' for d in es1_dns])
    es1_ip_str = "\n".join([f'      - "{i}"' for i in es1_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:
{es1_dns_str}
    ip:
{es1_ip_str}
"""
    if es2_host:
        if es2_dns is None:
            es2_dns = ["localhost", "es-node2", es2_host]
        if es2_ips is None:
            es2_ips = ["127.0.0.1", es2_host]
        es2_dns_str = "\n".join([f'      - "{d}"' for d in es2_dns])
        es2_ip_str = "\n".join([f'      - "{i}"' for i in es2_ips])
        content += f"""  - name: "elasticsearch-node2"
    dns:
{es2_dns_str}
    ip:
{es2_ip_str}
"""

    content += f"""  - 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


# ==============================================================================
# CERTIFICATE GENERATION CORE
# ==============================================================================
def generate_all_certificates(cert_dir, es1_dns, es1_ips, es2_host=None, es2_dns=None, es2_ips=None, kib_dns=None, kib_ips=None):
    ca_path = os.path.join(cert_dir, "elastic-stack-ca.p12")
    es_certutil = "/usr/share/elasticsearch/bin/elasticsearch-certutil"

    log_step("Ensuring certificate directory exists...")
    os.makedirs(cert_dir, exist_ok=True)

    log_step("Checking / Generating 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:
            ca_valid = True

    if not ca_valid:
        if os.path.exists(ca_path):
            os.remove(ca_path)
        run_cmd([es_certutil, "ca", "--out", ca_path, "--pass", ""])
        log_success("Fresh CA generated.")
    else:
        log_success("Valid passwordless CA found.")

    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_step("Generating node certificates via instances.yml...")
    instances_yml_path = os.path.join(cert_dir, "instances.yml")
    write_instances_yaml(instances_yml_path, es1_dns, es1_ips, es2_host, es2_dns, es2_ips, kib_dns, kib_ips)

    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", ""])

    if os.path.exists(certs_zip_path):
        with zipfile.ZipFile(certs_zip_path, "r") as zip_ref:
            zip_ref.extractall(cert_dir)

    # Node 1 ES certs
    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)

    # Node 2 ES certs (if present)
    es2_p12_path = os.path.join(cert_dir, "es2_node.p12")
    if es2_host:
        es2_p12_extracted = os.path.join(cert_dir, "elasticsearch-node2", "elasticsearch-node2.p12")
        if os.path.exists(es2_p12_extracted):
            shutil.move(es2_p12_extracted, es2_p12_path)

    # Kibana certs
    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 item in [os.path.join(cert_dir, "elasticsearch"), os.path.join(cert_dir, "elasticsearch-node2"), os.path.join(cert_dir, "kibana"), certs_zip_path, instances_yml_path, kibana_key_tmp]:
        if os.path.exists(item):
            if os.path.isdir(item):
                shutil.rmtree(item)
            else:
                os.remove(item)
    log_success("All node certificates generated and formatted.")

    return http_ca_temp, http_p12_path, transport_p12_path, es2_p12_path, kibana_crt_path, kibana_key_path


def clean_es_yaml_str(content):
    lines = content.splitlines()
    new_lines = []
    in_ignore_block = False
    ignore_indent = 0

    for line in lines:
        stripped = line.strip()
        if stripped.startswith("#"):
            new_lines.append(line)
            continue

        if "truststore" in stripped.lower() or "keystore" in stripped.lower():
            in_ignore_block = True
            ignore_indent = len(line) - len(line.lstrip())
            new_lines.append(f"# {line}")
            continue

        if in_ignore_block:
            current_indent = len(line) - len(line.lstrip())
            if stripped and current_indent > ignore_indent:
                new_lines.append(f"# {line}")
                continue
            else:
                in_ignore_block = False

        if re.search(r"(password|secure_key_passphrase)", stripped):
            new_lines.append(f"# {line}")
            continue

        new_lines.append(line)

    # Check active (non-comment) lines only
    active_text = "\n".join([l for l in new_lines if not l.strip().startswith("#")])

    result = "\n".join(new_lines) + "\n\n# Auto-generated clean SSL configuration by Autopilot\n"

    if "xpack.security.http.ssl.keystore.path" not in active_text:
        result += 'xpack.security.http.ssl.keystore.path: "/etc/elasticsearch/certs/http.p12"\n'
    if "xpack.security.transport.ssl.keystore.path" not in active_text:
        result += 'xpack.security.transport.ssl.keystore.path: "/etc/elasticsearch/certs/transport.p12"\n'
    if "xpack.security.http.ssl.certificate_authorities" not in active_text:
        result += 'xpack.security.http.ssl.certificate_authorities: ["/etc/elasticsearch/certs/http_ca.crt"]\n'
    if "xpack.security.transport.ssl.certificate_authorities" not in active_text:
        result += 'xpack.security.transport.ssl.certificate_authorities: ["/etc/elasticsearch/certs/http_ca.crt"]\n'

    return result


def sanitize_es_config_and_keystore(es_config_path, ssh_prefix=None):
    keys = [
        "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",
        "xpack.security.http.ssl.secure_key_passphrase",
        "xpack.security.transport.ssl.secure_key_passphrase",
    ]

    py_cleaner = r"""import re, sys
path = sys.argv[1]
try:
    c = open(path).read()
    lines = c.splitlines()
    n = []
    i_block = False
    i_indent = 0
    for l in lines:
        s = l.strip()
        if s.startswith("#"):
            n.append(l)
            continue
        if "truststore" in s.lower() or "keystore" in s.lower():
            i_block = True
            i_indent = len(l) - len(l.lstrip())
            n.append(f"# {l}")
            continue
        if i_block:
            curr = len(l) - len(l.lstrip())
            if s and curr > i_indent:
                n.append(f"# {l}")
                continue
            else:
                i_block = False
        if re.search(r"(password|secure_key_passphrase)", s):
            n.append(f"# {l}")
            continue
        n.append(l)
    active = '\n'.join([l for l in n if not l.strip().startswith('#')])
    res = '\n'.join(n) + '\n\n# Auto-generated clean SSL configuration by Autopilot\n'
    if 'xpack.security.http.ssl.keystore.path' not in active:
        res += 'xpack.security.http.ssl.keystore.path: "/etc/elasticsearch/certs/http.p12"\n'
    if 'xpack.security.transport.ssl.keystore.path' not in active:
        res += 'xpack.security.transport.ssl.keystore.path: "/etc/elasticsearch/certs/transport.p12"\n'
    if 'xpack.security.http.ssl.certificate_authorities' not in active:
        res += 'xpack.security.http.ssl.certificate_authorities: ["/etc/elasticsearch/certs/http_ca.crt"]\n'
    if 'xpack.security.transport.ssl.certificate_authorities' not in active:
        res += 'xpack.security.transport.ssl.certificate_authorities: ["/etc/elasticsearch/certs/http_ca.crt"]\n'
    open(path, 'w').write(res)
except Exception as e:
    pass
"""

    if ssh_prefix:
        remote_cmds = []
        remote_cmds.append(f"python3 -c {json.dumps(py_cleaner)} {es_config_path} 2>/dev/null || true")
        for k in keys:
            remote_cmds.append(f"/usr/share/elasticsearch/bin/elasticsearch-keystore remove {k} 2>/dev/null || true")
            remote_cmds.append(f"echo -n '' | /usr/share/elasticsearch/bin/elasticsearch-keystore add -x -f {k} 2>/dev/null || true")
        run_cmd(ssh_prefix + ["; ".join(remote_cmds)])
    else:
        if os.path.exists(es_config_path):
            try:
                with open(es_config_path, "r", encoding="utf-8") as f:
                    content = f.read()
                cleaned = clean_es_yaml_str(content)
                with open(es_config_path, "w", encoding="utf-8") as f:
                    f.write(cleaned)
            except Exception:
                pass

        es_keystore = "/usr/share/elasticsearch/bin/elasticsearch-keystore"
        if os.path.exists(es_keystore):
            for k in keys:
                run_cmd([es_keystore, "remove", k])
                run_cmd(["sh", "-c", f"echo -n '' | {es_keystore} add -x -f {k} 2>/dev/null || true"])


# ==============================================================================
# NODE 2 ELASTICSEARCH DISTRIBUTOR & RESTARTER
# ==============================================================================
def deploy_to_es2_node(cert_dir, es2_host, es2_user, es2_port, ssh_key, es2_p12_path, http_ca_temp):
    log_header(f"AUTOPILOT: Deploying Certificates to ES Node 2 ({es2_host})")

    ssh_base = build_ssh_prefix(es2_user, es2_host, es2_port, ssh_key)
    scp_base = build_scp_prefix(es2_port, ssh_key)
    remote_target = f"{es2_user}@{es2_host}" if es2_user else es2_host

    log_step(f"Ensuring remote cert directory exists on ES Node 2 ({es2_host})...")
    run_cmd(ssh_base + ["mkdir -p /etc/elasticsearch/certs /tmp/es2-certs-staging"])

    log_step("SCP node certificates to ES Node 2...")
    run_cmd(scp_base + [es2_p12_path, f"{remote_target}:/tmp/es2-certs-staging/http.p12"])
    run_cmd(scp_base + [http_ca_temp, f"{remote_target}:/tmp/es2-certs-staging/http_ca.crt"])

    log_step("Sanitizing remote ES Node 2 config & keystore passwords...")
    sanitize_es_config_and_keystore("/etc/elasticsearch/elasticsearch.yml", ssh_prefix=ssh_base)

    install_cmd = (
        "cp /tmp/es2-certs-staging/http.p12 /etc/elasticsearch/certs/http.p12 && "
        "cp /tmp/es2-certs-staging/http.p12 /etc/elasticsearch/certs/transport.p12 && "
        "cp /tmp/es2-certs-staging/http_ca.crt /etc/elasticsearch/certs/http_ca.crt && "
        "chown -R root:elasticsearch /etc/elasticsearch/certs && "
        "chmod 755 /etc/elasticsearch/certs && "
        "chmod 640 /etc/elasticsearch/certs/http.p12 /etc/elasticsearch/certs/transport.p12 && "
        "chmod 644 /etc/elasticsearch/certs/http_ca.crt && "
        "rm -rf /tmp/es2-certs-staging"
    )
    run_cmd(ssh_base + [install_cmd])
    log_success("Certificates installed & keystore passwords sanitized on ES Node 2.")

    log_step("Restarting Elasticsearch service on ES Node 2...")
    run_cmd(ssh_base + ["systemctl restart elasticsearch"])
    if not spin_until(f"Waiting for ES Node 2 ({es2_host}:9200)...", lambda: run_cmd(ssh_base + ["curl -k -s https://localhost:9200"]).returncode == 0, max_attempts=30, delay=6):
        log_error(f"ES Node 2 ({es2_host}) failed to start after 180s timeout!")
        log_warn(f"Logs excerpt from ES Node 2 ({es2_host}):")
        run_cmd(ssh_base + ["journalctl -u elasticsearch -n 25 --no-pager"], capture=False)
        sys.exit(1)
    log_success("ES Node 2 restarted and healthy.")


# ==============================================================================
# KIBANA CA & CONFIG SYNCHRONIZER
# ==============================================================================
def sync_kibana_ca_and_config(kibana_config_path, ca_src_path, new_fingerprint, ssh_prefix=None):
    py_kib_cleaner = r"""import re, sys, os, shutil
kib_config = sys.argv[1]
ca_src = sys.argv[2]
fingerprint = sys.argv[3]
ca_targets = [
    "/etc/kibana/certs/http_ca.crt",
    "/etc/kibana/certs/ca.crt",
    "/etc/kibana/certs/elasticsearch-ca.pem",
    "/var/lib/kibana/ca_1721297493415.crt"
]
try:
    if os.path.exists(kib_config):
        content = open(kib_config).read()
        m = re.search(r'certificateAuthorities:\s*(?:\[\s*)?["\']?([^"\',\s\]]+)["\']?', content)
        if m:
            ca_targets.append(m.group(1))
        for t in set(ca_targets):
            os.makedirs(os.path.dirname(t), exist_ok=True)
            shutil.copy(ca_src, t)
            os.chmod(t, 0o644)
        new_content = content
        if fingerprint:
            new_content = re.sub(r'(ca_trusted_fingerprint:\s*)[a-fA-F0-9]+', rf'\g<1>{fingerprint}', new_content)
            new_content = re.sub(r'(elasticsearch\.ssl\.ca_trusted_fingerprint:\s*)[a-fA-F0-9]+', rf'\g<1>{fingerprint}', new_content)
        new_content = re.sub(r'elasticsearch\.ssl\.certificateAuthorities:\s*["\']?([^"\',\s\]]+)["\']?(?!\s*\])', r'elasticsearch.ssl.certificateAuthorities: [ "\1" ]', new_content)
        if "certificateAuthorities" not in new_content:
            new_content += '\nelasticsearch.ssl.certificateAuthorities: [ "/etc/kibana/certs/http_ca.crt" ]\n'
        if new_content != content:
            open(kib_config, "w").write(new_content)
except Exception as e:
    pass
"""

    if ssh_prefix:
        remote_cmds = []
        remote_cmds.append(f"python3 -c {json.dumps(py_kib_cleaner)} {kibana_config_path} /tmp/kib-staging/http_ca.crt {new_fingerprint} 2>/dev/null || true")
        run_cmd(ssh_prefix + ["; ".join(remote_cmds)])
    else:
        if os.path.exists(kibana_config_path):
            try:
                content = open(kibana_config_path).read()
                ca_targets = [
                    "/etc/kibana/certs/http_ca.crt",
                    "/etc/kibana/certs/ca.crt",
                    "/etc/kibana/certs/elasticsearch-ca.pem",
                    "/var/lib/kibana/ca_1721297493415.crt"
                ]
                m = re.search(r'certificateAuthorities:\s*(?:\[\s*)?["\']?([^"\',\s\]]+)["\']?', content)
                if m:
                    ca_targets.append(m.group(1))
                for t in set(ca_targets):
                    os.makedirs(os.path.dirname(t), exist_ok=True)
                    shutil.copy(ca_src_path, t)
                    os.chmod(t, 0o644)
                new_content = content
                if new_fingerprint:
                    new_content = re.sub(r'(ca_trusted_fingerprint:\s*)[a-fA-F0-9]+', rf'\g<1>{new_fingerprint}', new_content)
                    new_content = re.sub(r'(elasticsearch\.ssl\.ca_trusted_fingerprint:\s*)[a-fA-F0-9]+', rf'\g<1>{new_fingerprint}', new_content)
                new_content = re.sub(r'elasticsearch\.ssl\.certificateAuthorities:\s*["\']?([^"\',\s\]]+)["\']?(?!\s*\])', r'elasticsearch.ssl.certificateAuthorities: [ "\1" ]', new_content)
                if "certificateAuthorities" not in new_content:
                    new_content += '\nelasticsearch.ssl.certificateAuthorities: [ "/etc/kibana/certs/http_ca.crt" ]\n'
                if new_content != content:
                    open(kibana_config_path, "w").write(new_content)
            except Exception:
                pass


# ==============================================================================
# AUTOPILOT ORCHESTRATOR
# ==============================================================================
def main():
    parser = argparse.ArgumentParser(description="Autopilot Universal Elastic Stack Certificate Renewal Script")
    parser.add_argument("--es2-host", default=os.getenv("ES2_HOST", ""), help="2nd Elasticsearch node IP/hostname")
    parser.add_argument("--es2-user", default=os.getenv("ES2_SSH_USER", "root"), help="SSH user for 2nd ES node")
    parser.add_argument("--es2-port", default=os.getenv("ES2_SSH_PORT", "22"), help="SSH port for 2nd ES node")

    parser.add_argument("--kibana-host", default=os.getenv("KIBANA_HOST", ""), help="Remote Kibana host IP")
    parser.add_argument("--kibana-user", default=os.getenv("KIBANA_SSH_USER", "root"), help="SSH user for Kibana host")
    parser.add_argument("--kibana-port", default=os.getenv("KIBANA_SSH_PORT", "22"), help="SSH port for Kibana host")

    parser.add_argument("--ssh-key", default=os.getenv("KIBANA_SSH_KEY", os.getenv("ES2_SSH_KEY", "")), help="Path to SSH private key")
    parser.add_argument("--es-config", default="/etc/elasticsearch/elasticsearch.yml", help="Path to elasticsearch.yml")
    parser.add_argument("--kibana-config", default="/etc/kibana/kibana.yml", help="Path to kibana.yml")

    args = parser.parse_args()

    if os.geteuid() != 0:
        print(f"{BOLD_RED}ERROR: Autopilot requires root privileges. Run with sudo.{RESET}", file=sys.stderr)
        sys.exit(1)

    cert_dir = os.getenv("CERT_DIR", "/etc/elasticsearch/certs")
    es_data = parse_simple_yaml(args.es_config)
    kib_data = parse_simple_yaml(args.kibana_config)

    es1_dns, es1_ips = discover_network_sans()
    kib_dns = ["localhost", "kibana-server", args.kibana_host] if args.kibana_host else es1_dns
    kib_ips = ["127.0.0.1", args.kibana_host] if args.kibana_host else es1_ips

    log_header("AUTOPILOT: Universal Elastic Stack Certificate Exchange")

    # Step 1: Generate All Node Certificates
    http_ca_temp, http_p12_path, transport_p12_path, es2_p12_path, kibana_crt_path, kibana_key_path = generate_all_certificates(
        cert_dir, es1_dns, es1_ips, args.es2_host, None, None, kib_dns, kib_ips
    )

    # Step 2: Apply to Local ES Node 1
    log_step("Applying certificates to Local ES Node 1...")
    es_http_p12_target = os.getenv("ES_HTTP_P12_TARGET", es_data.get("xpack.security.http.ssl.keystore.path", "/etc/elasticsearch/certs/http.p12"))
    es_transport_p12_target = os.getenv("ES_TRANSPORT_P12_TARGET", es_data.get("xpack.security.transport.ssl.keystore.path", "/etc/elasticsearch/certs/transport.p12"))
    es_http_ca_target = os.getenv("ES_HTTP_CA_TARGET", "/etc/elasticsearch/certs/http_ca.crt")

    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(["chown", "-R", "root:elasticsearch", cert_dir])
    run_cmd(["chmod", "755", cert_dir])
    for f in [http_p12_path, transport_p12_path]:
        if os.path.exists(f):
            run_cmd(["chmod", "640", f])

    log_step("Sanitizing Local ES Node 1 config & keystore passwords...")
    sanitize_es_config_and_keystore(args.es_config)

    run_cmd(["systemctl", "restart", "elasticsearch"])
    if not spin_until("Waiting for Local ES Node 1 on port 9200...", lambda: run_cmd(["curl", "-k", "-s", "https://localhost:9200"]).returncode == 0, max_attempts=30, delay=6):
        log_error("Local ES Node 1 failed to start after 180s timeout!")
        log_warn("Logs excerpt from Local ES Node 1:")
        run_cmd(["journalctl", "-u", "elasticsearch", "-n", "25", "--no-pager"], capture=False)
        sys.exit(1)

    # Step 3: Deploy to ES Node 2 (if specified)
    if args.es2_host:
        deploy_to_es2_node(cert_dir, args.es2_host, args.es2_user, args.es2_port, args.ssh_key, es2_p12_path, http_ca_temp)

    # Step 4: Deploy to Kibana & Fleet Server
    kib_crt_target = os.getenv("KIBANA_CRT_TARGET", kib_data.get("server.ssl.certificate", "/etc/kibana/certs/kibana.crt"))
    kib_key_target = os.getenv("KIBANA_KEY_TARGET", kib_data.get("server.ssl.key", "/etc/kibana/certs/kibana.key"))
    kib_ca_target = os.getenv("KIBANA_CA_TARGET", "/var/lib/kibana/ca_1721297493415.crt")

    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()

    if args.kibana_host:
        log_header(f"AUTOPILOT: Deploying Certificates to Remote Kibana Host ({args.kibana_host})")
        ssh_base = build_ssh_prefix(args.kibana_user, args.kibana_host, args.kibana_port, args.ssh_key)
        scp_base = build_scp_prefix(args.kibana_port, args.ssh_key)
        remote_target = f"{args.kibana_user}@{args.kibana_host}" if args.kibana_user else args.kibana_host

        remote_crt_dir = os.path.dirname(kib_crt_target)
        run_cmd(ssh_base + [f"mkdir -p {remote_crt_dir} /var/lib/kibana /tmp/kib-staging"])
        run_cmd(scp_base + [kibana_crt_path, f"{remote_target}:/tmp/kib-staging/kibana.crt"])
        run_cmd(scp_base + [kibana_key_path, f"{remote_target}:/tmp/kib-staging/kibana.key"])
        run_cmd(scp_base + [http_ca_temp, f"{remote_target}:/tmp/kib-staging/http_ca.crt"])

        remote_install_cmd = (
            f"cp /tmp/kib-staging/kibana.crt {kib_crt_target} && "
            f"cp /tmp/kib-staging/kibana.key {kib_key_target} && "
            f"cp /tmp/kib-staging/http_ca.crt {kib_ca_target} && "
            f"chown -R root:kibana {remote_crt_dir} && "
            f"chmod 644 {kib_crt_target} && "
            f"chmod 640 {kib_key_target} && "
            f"chmod 644 {kib_ca_target} && "
            f"rm -rf /tmp/kib-staging"
        )
        run_cmd(ssh_base + [remote_install_cmd])

        sync_kibana_ca_and_config(args.kibana_config, http_ca_temp, new_fingerprint, ssh_prefix=ssh_base)

        run_cmd(ssh_base + ["systemctl restart kibana"])
        spin_until(f"Waiting for remote Kibana on {args.kibana_host}:5601...", lambda: run_cmd(ssh_base + ["curl -k -s https://localhost:5601"]).returncode == 0)

        log_step("Generating Fleet Server token & re-enrolling remote Agent...")
        token_name = f"autopilot-{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)

        fleet_url = os.getenv("FLEET_SERVER_URL", f"https://{args.kibana_host}:8220")
        fleet_es_url = os.getenv("FLEET_SERVER_ES_URL", "https://10.1.10.3:9200")
        run_cmd(ssh_base + ["systemctl stop elastic-agent || true"])
        remote_enroll_cmd = (
            f"/opt/Elastic/Agent/elastic-agent enroll --force "
            f"--url={fleet_url} "
            f"--fleet-server-es={fleet_es_url} "
            f"--fleet-server-es-ca={kib_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=8220 "
            f"--insecure"
        )
        run_cmd(ssh_base + [remote_enroll_cmd])
        run_cmd(ssh_base + ["systemctl start elastic-agent"])
        spin_until(f"Waiting for remote Fleet Server on {args.kibana_host}:8220...", lambda: run_cmd(ssh_base + ["curl -k -s https://localhost:8220"]).returncode == 0)

    else:
        log_step("Applying Kibana certificates locally...")
        symlink_cert(kibana_crt_path, kib_crt_target)
        symlink_cert(kibana_key_path, kib_key_target)
        symlink_cert(http_ca_temp, kib_ca_target)

        for f in [kibana_crt_path, kibana_key_path]:
            if os.path.exists(f):
                run_cmd(["chown", "root:kibana", f])
        if os.path.exists(kibana_crt_path):
            run_cmd(["chmod", "644", kibana_crt_path])
        if os.path.exists(kibana_key_path):
            run_cmd(["chmod", "640", kibana_key_path])

        run_cmd(["usermod", "-aG", "elasticsearch", "kibana"])

        sync_kibana_ca_and_config(args.kibana_config, http_ca_temp, new_fingerprint)

        run_cmd(["systemctl", "restart", "kibana"])
        spin_until("Waiting for Kibana on port 5601...", lambda: run_cmd(["curl", "-k", "-s", "https://localhost:5601"]).returncode == 0)
        token_name = f"autopilot-{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)

        fleet_es_url = os.getenv("FLEET_SERVER_ES_URL", kib_data.get("elasticsearch.hosts", "https://10.1.10.3:9200"))
        fleet_url = os.getenv("FLEET_SERVER_URL", kib_data.get("fleet.agent.fleet_server.hosts", "https://10.1.10.3:8220"))
        run_cmd(["systemctl", "stop", "elastic-agent"])
        enroll_cmd = [
            "/opt/Elastic/Agent/elastic-agent", "enroll", "--force",
            f"--url={fleet_url}",
            f"--fleet-server-es={fleet_es_url}",
            f"--fleet-server-es-ca={es_http_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=8220",
            "--insecure"
        ]
        run_cmd(enroll_cmd)
        run_cmd(["systemctl", "start", "elastic-agent"])
        spin_until("Waiting for Fleet Server endpoint on port 8220...", lambda: run_cmd(["curl", "-k", "-s", "https://localhost:8220"]).returncode == 0)

    log_header("✔ ALL-IN-ONE AUTOPILOT COMPLETED SUCCESSFULLY!")


if __name__ == "__main__":
    main()