#!/usr/bin/env python3
"""
==============================================================================
Elastic Stack Certificate Renewal & Service Sync (Python)
==============================================================================
Python implementation of update-certs.sh. Generates/renews CA and node
certificates, configures Elasticsearch and Kibana permissions/symlinks, updates
CA fingerprints, restarts stack services, and re-enrolls Fleet Server 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

# 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, san_dns, san_ips):
    dns_entries = "\n".join([f'      - "{d}"' for d in san_dns])
    ip_entries = "\n".join([f'      - "{i}"' for i in san_ips])

    content = f"""instances:
  - name: "elasticsearch"
    dns:
{dns_entries}
    ip:
{ip_entries}
  - name: "kibana"
    dns:
{dns_entries}
    ip:
{ip_entries}
"""
    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 main():
    if os.geteuid() != 0:
        print(f"{BOLD_RED}ERROR: Please run this script as root (using sudo){RESET}", file=sys.stderr)
        sys.exit(1)

    # Environment & Default Configuration
    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")

    kibana_crt_target = os.getenv("KIBANA_CRT_TARGET", "/etc/kibana/certs/kibana.crt")
    kibana_key_target = os.getenv("KIBANA_KEY_TARGET", "/etc/kibana/certs/kibana.key")
    kibana_ca_target = os.getenv("KIBANA_CA_TARGET", "/var/lib/kibana/ca_1721297493415.crt")

    kibana_config = os.getenv("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", "https://10.1.10.3:8220")
    fleet_server_policy = os.getenv("FLEET_SERVER_POLICY", "fleet-server-policy")
    fleet_server_port = os.getenv("FLEET_SERVER_PORT", "8220")

    san_dns_str = os.getenv("SAN_DNS", "localhost cmsrv2024 siem2024")
    san_ips_str = os.getenv("SAN_IPS", "127.0.0.1 10.1.10.3 192.168.196.95")

    san_dns = [d for d in san_dns_str.strip().split() if d]
    san_ips = [i for i in san_ips_str.strip().split() if i]

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

    log_section("Elastic Stack Certificate Renewal & Service Sync (Python)")

    log_step("Ensuring certificate directory exists...")
    os.makedirs(cert_dir, exist_ok=True)
    log_sub(f"Target directory: {cert_dir}")

    log_step("Checking Certificate Authority (CA) status...")
    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.")
    else:
        log_sub("No existing CA found.")

    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 instances configuration...")
    instances_yml_path = os.path.join(cert_dir, "instances.yml")
    write_instances_yaml(instances_yml_path, san_dns, san_ips)
    log_sub("instances.yml written.")

    log_step("Generating node certificates...")
    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 and organizing certificate files...")
    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_success("Certificates converted to required formats.")

    log_step("Clearing legacy secure keystore passwords...")
    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_sub("Keystore parameters updated.")

    log_step("Applying strict file ownership and permissions...")
    run_cmd(["chown", "root:elasticsearch", cert_dir])
    run_cmd(["chmod", "755", cert_dir])

    for f in [ca_path, http_p12_path, transport_p12_path]:
        if os.path.exists(f):
            run_cmd(["chown", "root:elasticsearch", f])
            run_cmd(["chmod", "640", f])

    if os.path.exists(http_ca_temp):
        run_cmd(["chown", "root:elasticsearch", http_ca_temp])
        run_cmd(["chmod", "644", http_ca_temp])

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

    service_tokens = "/etc/elasticsearch/service_tokens"
    if os.path.exists(service_tokens):
        run_cmd(["chown", "root:elasticsearch", service_tokens])
        run_cmd(["chmod", "660", service_tokens])

    run_cmd(["usermod", "-aG", "elasticsearch", "kibana"])
    log_success("Permissions set.")

    log_step("Updating symlinks to target service locations...")
    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)
    symlink_cert(kibana_crt_path, kibana_crt_target)
    symlink_cert(kibana_key_path, kibana_key_target)
    symlink_cert(http_ca_temp, kibana_ca_target)

    for target in [es_http_p12_target, es_transport_p12_target, es_http_ca_target]:
        if os.path.exists(target):
            run_cmd(["chown", "root:elasticsearch", target])
    if os.path.exists(kibana_crt_target):
        run_cmd(["chown", "-R", "root:kibana", os.path.dirname(kibana_crt_target)])
    if os.path.exists(kibana_ca_target):
        run_cmd(["chown", "kibana:kibana", kibana_ca_target])
    log_success("Symlinks updated.")

    log_step("Computing new CA fingerprint & updating Kibana...")
    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 os.path.exists(kibana_config) and new_fingerprint:
        with open(kibana_config, "r", encoding="utf-8") as f:
            k_content = f.read()
        k_content = re.sub(r"ca_trusted_fingerprint:\s*[a-fA-F0-9]+", f"ca_trusted_fingerprint: {new_fingerprint}", k_content)
        with open(kibana_config, "w", encoding="utf-8") as f:
            f.write(k_content)
        log_success("Kibana config updated.")

    log_section("Restarting Core Stack Services")

    log_step("Restarting Elasticsearch service...")
    run_cmd(["systemctl", "restart", "elasticsearch"])

    def check_es():
        res = run_cmd(["curl", "-k", "-s", "https://localhost:9200"])
        return res.returncode == 0

    if not spin_until("Waiting for Elasticsearch to start (port 9200)...", check_es, max_attempts=30, delay=6):
        log_warn("Elasticsearch logs excerpt:")
        run_cmd(["journalctl", "-u", "elasticsearch", "-n", "15", "--no-pager"], capture=False)

    log_step("Restarting Kibana service...")
    run_cmd(["systemctl", "restart", "kibana"])

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

    if not spin_until("Waiting for Kibana to start (port 5601)...", check_kibana):
        log_warn("Kibana logs excerpt:")
        run_cmd(["journalctl", "-u", "kibana", "-n", "15", "--no-pager"], capture=False)

    log_section("Re-enrolling Elastic Agent (Fleet Server)")

    es_sa_token = ""
    if os.path.exists(kibana_config):
        with open(kibana_config, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip().startswith("elasticsearch.serviceAccountToken:"):
                    es_sa_token = line.split(":", 1)[1].strip().strip('"')

    pre_agent_ids = []
    hostname = socket.gethostname()
    if es_sa_token:
        try:
            url = "https://localhost:9200/.fleet-agents/_search"
            data = json.dumps({"query": {"term": {"local_metadata.host.hostname": hostname}}}).encode("utf-8")
            req = urllib.request.Request(url, data=data, headers={"Authorization": f"Bearer {es_sa_token}", "Content-Type": "application/json"}, method="POST")
            ctx = ssl._create_unverified_context()
            with urllib.request.urlopen(req, context=ctx, timeout=5) as resp:
                body = json.loads(resp.read().decode("utf-8"))
                hits = body.get("hits", {}).get("hits", [])
                pre_agent_ids = [h["_id"] for h in hits]
            log_sub(f"Captured {len(pre_agent_ids)} pre-existing agent ID(s) for host {BOLD_CYAN}{hostname}{RESET}.")
        except Exception:
            pass

    log_step("Preparing Fleet Server service token before touching Agent...")
    token_name = f"cert-exchange-{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():
            parts = res.stdout.strip().split()
            fleet_service_token = parts[-1]
            break
        log_sub(f"Attempt {attempt}/12: Elasticsearch initializing token store, waiting 5s...")
        time.sleep(5)

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

    if os.path.exists(service_tokens):
        run_cmd(["chown", "root:elasticsearch", service_tokens])
        run_cmd(["chmod", "660", service_tokens])

    log_step("Stopping Elastic Agent service...")
    run_cmd(["systemctl", "stop", "elastic-agent"])
    log_success("Elastic Agent stopped.")

    log_step("Re-enrolling Elastic Agent (Fleet Server)...")
    enroll_cmd = [
        "/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={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={fleet_server_port}",
        "--insecure"
    ]
    run_cmd(enroll_cmd)
    log_success("Elastic Agent re-enrolled with new CA fingerprint.")

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

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

    spin_until("Waiting for Fleet Server endpoint (port 8220)...", check_fleet)

    if pre_agent_ids and es_sa_token:
        log_step("Cleaning up former Agent entry/entries from Kibana Fleet UI...")
        try:
            url = "https://localhost:9200/.fleet-agents/_search"
            data = json.dumps({
                "query": {"term": {"local_metadata.host.hostname": hostname}},
                "sort": [{"last_checkin": {"order": "desc"}}]
            }).encode("utf-8")
            req = urllib.request.Request(url, data=data, headers={"Authorization": f"Bearer {es_sa_token}", "Content-Type": "application/json"}, method="POST")
            ctx = ssl._create_unverified_context()
            new_agent_id = ""
            with urllib.request.urlopen(req, context=ctx, timeout=5) as resp:
                body = json.loads(resp.read().decode("utf-8"))
                hits = body.get("hits", {}).get("hits", [])
                if hits:
                    new_agent_id = hits[0]["_id"]

            for old_id in pre_agent_ids:
                if old_id != new_agent_id:
                    del_url = f"https://localhost:9200/.fleet-agents/_doc/{old_id}"
                    del_req = urllib.request.Request(del_url, headers={"Authorization": f"Bearer {es_sa_token}"}, method="DELETE")
                    with urllib.request.urlopen(del_req, context=ctx, timeout=5):
                        pass
                    log_sub(f"Removed obsolete agent ID: {DIM}{old_id}{RESET}")
        except Exception:
            pass
        log_success("Duplicate/stale agent entries purged from Fleet UI.")

    log_section("Final Service & Certificate Verification")

    log_step("Elasticsearch (10.1.10.3:9200):")
    res1 = run_cmd(["openssl", "s_client", "-connect", "10.1.10.3:9200", "-showcerts"], text=True)
    if res1.returncode == 0 and "END CERTIFICATE" in res1.stdout:
        log_success("Elasticsearch SSL endpoint verified.")
    else:
        log_error("Failed to verify Elasticsearch SSL endpoint.")

    log_step("Kibana (10.1.10.3:5601):")
    res2 = run_cmd(["openssl", "s_client", "-connect", "10.1.10.3:5601", "-showcerts"], text=True)
    if res2.returncode == 0 and "END CERTIFICATE" in res2.stdout:
        log_success("Kibana SSL endpoint verified.")
    else:
        log_error("Failed to verify Kibana SSL endpoint.")

    log_step("Fleet Server (10.1.10.3:8220):")
    res3 = run_cmd(["openssl", "s_client", "-connect", "10.1.10.3:8220", "-showcerts"], text=True)
    if res3.returncode == 0 and "END CERTIFICATE" in res3.stdout:
        log_success("Fleet Server SSL endpoint verified.")
    else:
        log_error("Failed to verify Fleet Server SSL endpoint.")

    print(f"\n{BOLD_GREEN}===================================================={RESET}")
    print(f" {BOLD_GREEN}✔ All certificates renewed and services active!{RESET}")
    print(f"{BOLD_GREEN}===================================================={RESET}\n")


if __name__ == "__main__":
    main()