import paramiko
import urllib.request
import ssl
import json
import socket
import concurrent.futures

key_path = r'C:\Users\Snarf\.ssh\id_ed25519_snarfnet'

def check_k8s_cluster(host, name):
    res = {'name': name, 'host': host, 'errors': [], 'warnings': []}
    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect(host, username='root', key_filename=key_path, timeout=10)

        # Check nodes
        stdin, stdout, stderr = ssh.exec_command("k3s kubectl get nodes --no-headers")
        nodes_out = stdout.read().decode().strip().splitlines()
        not_ready = [n for n in nodes_out if not " Ready" in n]
        res['total_nodes'] = len(nodes_out)
        res['not_ready_nodes'] = not_ready

        # Check failing pods across all namespaces
        stdin, stdout, stderr = ssh.exec_command("k3s kubectl get pods -A --field-selector=status.phase!=Running,status.phase!=Succeeded --no-headers")
        failing_pods = stdout.read().decode().strip().splitlines()
        res['failing_pods'] = failing_pods

        # Check DNS namespace pods
        stdin, stdout, stderr = ssh.exec_command("k3s kubectl get pods -n dns -o wide --no-headers")
        dns_pods = stdout.read().decode().strip().splitlines()
        res['dns_pods'] = dns_pods

        # Check DNS service VIP
        stdin, stdout, stderr = ssh.exec_command("k3s kubectl get svc -n dns -o wide --no-headers")
        dns_svc = stdout.read().decode().strip().splitlines()
        res['dns_svc'] = dns_svc

        ssh.close()
    except Exception as e:
        res['errors'].append(f"SSH failed: {e}")
    return res

def check_swarm():
    res = {'name': 'Docker Swarm / NPM', 'host': '140.44.4.71', 'errors': []}
    try:
        ssh = paramiko.SSHClient()
        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh.connect('140.44.4.71', username='snarf', key_filename=key_path, timeout=10)

        stdin, stdout, stderr = ssh.exec_command("docker node ls --format '{{.Hostname}}: {{.Status}} ({{.Availability}})'")
        nodes = stdout.read().decode().strip().splitlines()
        res['nodes'] = nodes

        stdin, stdout, stderr = ssh.exec_command("docker ps --format '{{.Names}}: {{.Status}}'")
        containers = stdout.read().decode().strip().splitlines()
        res['containers'] = containers

        ssh.close()
    except Exception as e:
        res['errors'].append(f"SSH failed: {e}")
    return res

def test_dns_resolution(server_ip, domain="google.com"):
    import struct
    try:
        # Construct raw DNS query packet
        # Transaction ID: 0x1234, Flags: standard query (0x0100), Questions: 1
        header = struct.pack(">HHHHHH", 0x1234, 0x0100, 1, 0, 0, 0)
        qname = b"".join(bytes([len(p)]) + p.encode('ascii') for p in domain.split(".")) + b"\x00"
        question = qname + struct.pack(">HH", 1, 1) # Type A, Class IN
        packet = header + question

        sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        sock.settimeout(3.0)
        sock.sendto(packet, (server_ip, 53))
        data, _ = sock.recvfrom(512)
        sock.close()
        # Check RCODE in response
        flags = struct.unpack(">H", data[2:4])[0]
        rcode = flags & 0x000F
        return {'status': 'OK', 'rcode': rcode, 'bytes': len(data)}
    except Exception as e:
        return {'status': 'FAIL', 'error': str(e)}

def check_endpoint(url):
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE

    req = urllib.request.Request(url, headers={'User-Agent': 'AntigravityAudit/2.0'})
    try:
        # Don't follow redirects to capture 302/200 properly
        class NoRedirectHandler(urllib.request.HTTPRedirectHandler):
            def http_error_302(self, req, fp, code, msg, headers):
                return fp
            http_error_301 = http_error_302
            http_error_303 = http_error_302
            http_error_307 = http_error_302
            http_error_308 = http_error_302

        opener = urllib.request.build_opener(NoRedirectHandler)
        with opener.open(req, timeout=5) as resp:
            code = resp.status if hasattr(resp, 'status') else resp.code
            location = resp.headers.get('Location', '')
            server = resp.headers.get('Server', '')
            return {'url': url, 'code': code, 'location': location, 'server': server, 'ok': code in [200, 301, 302, 307, 308]}
    except urllib.error.HTTPError as e:
        return {'url': url, 'code': e.code, 'error': str(e), 'ok': e.code in [200, 301, 302, 307, 308, 401, 403]}
    except Exception as e:
        return {'url': url, 'code': 0, 'error': str(e), 'ok': False}

print("=== STARTING FULL INFRASTRUCTURE & ENDPOINT AUDIT ===")

# 1. Check K8s Clusters
aero = check_k8s_cluster('140.44.4.11', 'Aerodrome')
space = check_k8s_cluster('140.44.4.41', 'Spaceport')
swarm = check_swarm()

# 2. Check DNS Port 53 on both VIPs
dns_aero = test_dns_resolution('140.44.4.36', 'google.com')
dns_space = test_dns_resolution('140.44.4.65', 'google.com')
dns_local = test_dns_resolution('140.44.4.36', 'authentik.snarfnet.net')

# 3. Check Web Endpoints
endpoints_to_check = [
    'https://authentik.snarfnet.net',
    'https://pihole1.snarfnet.net/admin/',
    'https://pihole2.snarfnet.net/admin/',
    'https://dns1.snarfnet.net',
    'https://dns2.snarfnet.net',
    'https://git.snarfnet.net',
    'https://vault.snarfnet.net',
    'https://rancher.snarfnet.net',
    'https://argocd.snarfnet.net',
    'https://zabbix.snarfnet.net',
    'https://semaphore.snarfnet.net',
    'https://kuma.snarfnet.net',
    'https://homepage.snarfnet.net',
    'https://grafana.snarfnet.net'
]

endpoint_results = []
with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
    endpoint_results = list(executor.map(check_endpoint, endpoints_to_check))

audit_data = {
    'aerodrome': aero,
    'spaceport': space,
    'swarm': swarm,
    'dns_vip_aerodrome': dns_aero,
    'dns_vip_spaceport': dns_space,
    'dns_vip_internal': dns_local,
    'endpoints': endpoint_results
}

print(json.dumps(audit_data, indent=2))