#!/usr/bin/env python3
# ===============================================================
# © Daniel Harding & RomanAILabs. All rights reserved.
# Credits: GPT-5 (OpenAI) and contributors
# Project: RomanAILabs 4.0 Quantum Traceroute
# ===============================================================

import os, socket, time, json
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor

try:
    import scapy.all as scapy
    import geoip2.database
except ImportError:
    print("[!] Install: pip install scapy geoip2")

# Load MaxMind GeoLite2 if available (download .mmdb file separately)
GEOIP_DB = "/usr/share/GeoIP/GeoLite2-City.mmdb"
ASN_DB = "/usr/share/GeoIP/GeoLite2-ASN.mmdb"

geo_reader = None
asn_reader = None
if os.path.exists(GEOIP_DB):
    geo_reader = geoip2.database.Reader(GEOIP_DB)
if os.path.exists(ASN_DB):
    asn_reader = geoip2.database.Reader(ASN_DB)

# ===============================================================
# 4D WRAPPER
# ===============================================================
def quantum4d(func):
    """Decorator to run any function in parallel 'dimensions'."""
    def wrapper(*args, **kwargs):
        start = time.time()
        with ThreadPoolExecutor(max_workers=3) as ex:
            futures = [ex.submit(func, *args, **kwargs) for _ in range(3)]
            results = [f.result() for f in futures]
        end = time.time()
        print(f"[4D] {func.__name__} executed in {end-start:.2f}s across 3 dimensions.\n")
        return results[0]  # return one consolidated result
    return wrapper

# ===============================================================
# GEO LOOKUP HELPERS
# ===============================================================
def lookup_geo(ip):
    data = {"ip": ip}
    try:
        if geo_reader:
            geo = geo_reader.city(ip)
            data.update({
                "country": geo.country.name,
                "city": geo.city.name,
                "lat": geo.location.latitude,
                "lon": geo.location.longitude
            })
        if asn_reader:
            asn = asn_reader.asn(ip)
            data.update({
                "asn": asn.autonomous_system_number,
                "isp": asn.autonomous_system_organization
            })
    except Exception:
        pass
    return data

# ===============================================================
# 4D TRACEROUTE FUNCTION
# ===============================================================
@quantum4d
def traceroute_plus(target, max_hops=30, timeout=2):
    print(f"\n[4D-TRACE] Starting advanced traceroute to {target}\n")
    hops = []

    for ttl in range(1, max_hops+1):
        # Multi-protocol probes
        pkt_icmp = scapy.IP(dst=target, ttl=ttl) / scapy.ICMP()
        pkt_tcp = scapy.IP(dst=target, ttl=ttl) / scapy.TCP(dport=80, flags="S")
        pkt_udp = scapy.IP(dst=target, ttl=ttl) / scapy.UDP(dport=33434)

        start_time = time.time()
        reply_icmp = scapy.sr1(pkt_icmp, verbose=0, timeout=timeout)
        reply_tcp  = scapy.sr1(pkt_tcp, verbose=0, timeout=timeout)
        reply_udp  = scapy.sr1(pkt_udp, verbose=0, timeout=timeout)
        rtt = (time.time() - start_time) * 1000

        reply = reply_icmp or reply_tcp or reply_udp
        if reply:
            ip = reply.src
            geo = lookup_geo(ip)
            hop_info = {
                "ttl": ttl,
                "ip": ip,
                "rtt_ms": round(rtt, 2),
                "timestamp": datetime.utcnow().isoformat(),
                "geo": geo
            }
            hops.append(hop_info)

            # Verbose output
            print(f"[{ttl}] {ip} | {geo.get('country','?')}, {geo.get('city','?')} "
                  f"| ASN {geo.get('asn','?')} {geo.get('isp','?')} "
                  f"| RTT {hop_info['rtt_ms']} ms")

            if ip == target:
                break
        else:
            print(f"[{ttl}] *")

    # Save JSON for map/analysis
    with open("4d_traceroute.json", "w") as f:
        json.dump(hops, f, indent=2)

    print("\n[4D-TRACE] Complete. Results saved to 4d_traceroute.json\n")
    return hops

# ===============================================================
# TEST MENU
# ===============================================================
if __name__ == "__main__":
    tgt = input("Enter target host/IP: ")
    traceroute_plus(tgt)