#!/usr/bin/env python3
"""
Deploy and configure authoritative Chrony NTP servers:
- Primary: aerodrome-ctl01 (140.44.4.11)
- Secondary: nutserver (14.4.4.44)

Includes:
- Exterior Stratum 1 & 2 sources (Cloudflare, Google, NIST, pool.ntp.org)
- Peer exchange between primary and secondary
- Subnet allowances for homelab VLANs
- Local stratum 10 fallback for offline WAN resilience
- Verification of UDP 123 responsiveness
"""

import os
import sys
import time
import socket
import paramiko

sys.stdout.reconfigure(encoding="utf-8")

PRIV_KEY_PATH = os.path.expanduser("~/.ssh/id_ed25519_snarfnet")

PRIMARY_CONF = """# Snarfnet Authoritative Internal NTP Server (Primary)
# Upstream Stratum 1 & 2 Sources
server time.cloudflare.com iburst minpoll 4 maxpoll 8
server time.google.com iburst minpoll 4 maxpoll 8
server time.nist.gov iburst minpoll 4 maxpoll 8
server 0.us.pool.ntp.org iburst minpoll 4 maxpoll 8
server 1.us.pool.ntp.org iburst minpoll 4 maxpoll 8

# Peer with Secondary NTP Server
peer 14.4.4.44 minpoll 4 maxpoll 8

# Drift tracking & step adjustment
driftfile /var/lib/chrony/chrony.drift
makestep 1.0 3
rtcsync

# Allowed Internal Subnets
allow 140.44.4.0/24
allow 14.4.4.0/24
allow 14.10.10.0/24
allow 140.44.44.0/24
allow 14.44.4.0/24
allow 127.0.0.1
allow ::1

# Local stratum 10 fallback if internet WAN is disconnected
local stratum 10

# Logging
logchange 0.5
logdir /var/log/chrony
"""

SECONDARY_CONF = """# Snarfnet Authoritative Internal NTP Server (Secondary)
# Upstream Stratum 1 & 2 Sources
server time.cloudflare.com iburst minpoll 4 maxpoll 8
server time.google.com iburst minpoll 4 maxpoll 8
server time.nist.gov iburst minpoll 4 maxpoll 8
server 0.us.pool.ntp.org iburst minpoll 4 maxpoll 8
server 1.us.pool.ntp.org iburst minpoll 4 maxpoll 8

# Peer with Primary NTP Server
peer 140.44.4.11 minpoll 4 maxpoll 8

# Drift tracking & step adjustment
driftfile /var/lib/chrony/chrony.drift
makestep 1.0 3
rtcsync

# Allowed Internal Subnets
allow 140.44.4.0/24
allow 14.4.4.0/24
allow 14.10.10.0/24
allow 140.44.44.0/24
allow 14.44.4.0/24
allow 127.0.0.1
allow ::1

# Local stratum 10 fallback if internet WAN is disconnected
local stratum 10

# Logging
logchange 0.5
logdir /var/log/chrony
"""

SERVERS = [
    ("aerodrome-ctl01 (Primary)", "140.44.4.11", "root", PRIMARY_CONF),
    ("nutserver (Secondary)",    "14.4.4.44",   "snarf", SECONDARY_CONF),
]

def check_udp_123(host):
    msg = b'\x1b' + 47 * b'\0'
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sock.settimeout(3.0)
    try:
        sock.sendto(msg, (host, 123))
        data, addr = sock.recvfrom(1024)
        return True, len(data)
    except Exception as e:
        return False, str(e)
    finally:
        sock.close()

def main():
    pkey = paramiko.Ed25519Key.from_private_key_file(PRIV_KEY_PATH)

    for name, ip, user, conf_content in SERVERS:
        print(f"\n==================================================")
        print(f"Deploying Chrony on {name} [{ip}] as {user}...")
        print(f"==================================================")
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            client.connect(ip, username=user, pkey=pkey, timeout=15)
            
            # 1. Stop systemd-timesyncd if present, update apt & install chrony
            install_cmd = """
            export DEBIAN_FRONTEND=noninteractive
            if [ "$(id -u)" -ne 0 ]; then SUDO="sudo"; else SUDO=""; fi
            $SUDO systemctl stop systemd-timesyncd 2>/dev/null || true
            $SUDO systemctl disable systemd-timesyncd 2>/dev/null || true
            $SUDO apt-get update -y
            $SUDO apt-get install -y chrony
            """
            print("Installing chrony package...")
            stdin, stdout, stderr = client.exec_command(install_cmd)
            stdout.channel.recv_exit_status()
            
            # 2. Write configuration
            print("Writing /etc/chrony/chrony.conf...")
            write_cmd = f"""
            if [ "$(id -u)" -ne 0 ]; then SUDO="sudo"; else SUDO=""; fi
            $SUDO mkdir -p /etc/chrony /var/lib/chrony /var/log/chrony
            $SUDO tee /etc/chrony/chrony.conf << 'EOF' > /dev/null
{conf_content}
EOF
            $SUDO systemctl unmask chrony 2>/dev/null || true
            $SUDO systemctl restart chrony
            $SUDO systemctl enable chrony
            """
            stdin, stdout, stderr = client.exec_command(write_cmd)
            stdout.channel.recv_exit_status()
            
            # 3. Wait briefly for chronyd to bind and acquire sources
            print("Waiting 5 seconds for chronyd stabilization...")
            time.sleep(5)
            
            # 4. Check chronyc sources and tracking
            verify_cmd = """
            if [ "$(id -u)" -ne 0 ]; then SUDO="sudo"; else SUDO=""; fi
            echo "--- Chrony Sources ---"
            $SUDO chronyc sources -v
            echo "--- Chrony Tracking ---"
            $SUDO chronyc tracking
            """
            stdin, stdout, stderr = client.exec_command(verify_cmd)
            out = stdout.read().decode('utf-8', errors='replace')
            print(out)
            
        except Exception as e:
            print(f"Error on {name}: {e}")
        finally:
            client.close()

    print("\n==================================================")
    print("Verifying UDP 123 Reachability from Network...")
    print("==================================================")
    for name, ip, _, _ in SERVERS:
        ok, res = check_udp_123(ip)
        print(f"NTP Query to {name} ({ip}:123) -> {'SUCCESS (' + str(res) + ' bytes returned)' if ok else 'FAILED: ' + str(res)}")

if __name__ == "__main__":
    main()