#!/usr/bin/env python3.11


import subprocess
import sqlite3
import re
import json
import ipaddress
from datetime import datetime

def run_traceroute(host):
    timestamp = datetime.now().timestamp()
    result = subprocess.run(['traceroute', host], stdout=subprocess.PIPE)
    return result.stdout.decode(), timestamp

def parse_traceroute_output(output, timestamp):

    lines = output.strip().split('\n')
    hops = []

    ip_regex = r"\((.*?)\)" # ipaddress are in ()

    target = output.strip().split('\n')[0].split()[2]

    for line in lines[1:]:
        hop = {}
        hop_info = line.split()
        hop_number = int(hop_info[0])
        hop_name = None
        hop_ip = None
        hop_latency = None

        latencies = []

        #print("##### "+str(hop_info))
        count = 0
        for part in hop_info[1:]:
            count += 1

            # source node drops or blocks icmp packages
            # We will give funny to name to hop for not answering and move on.
            if part == '*':
                hop_name = 'hop.'+str(count)+'.'+str(target)
                break


            # If first colum is either name or ip-address
            if count == 1:
                match = re.search(ip_regex, part)
                if  match:
                    hop_ip = part.strip('()')
                else:
                    hop_name = part

            # Second colum is ip-address first latency reading
            if count == 2:
                if re.search(ip_regex, part):
                    try:
                        _ip = ipaddress.ip_address(part.strip('()'))
                        hop_ip = part.strip('()')
                    except ValueError:
                        pass  # Ignore if it's not a valid IP address

            # Rest of the input colums are either latency floats, 'ms' or 
            # reruns of the hop_name and hop_ip...
            # We only need the latency floats anymore.
            else:
                try:
                    latency = float(part)
                    latencies.append(latency)
                except ValueError:
                    pass

        hop_latency = sum(latencies) / len(latencies) if latencies else None

        hop['timestamp'] = timestamp
        hop['hop_number'] = hop_number
        if not hop_name == None:
            hop['hop_name'] = hop_name
        hop['hop_ip'] = hop_ip
        hop['hop_latency'] = hop_latency
        
        hops.append(target)
        hops.append(hop)

    return hops


def create_tables(databasefile):
        # Connect to the SQLite database
    conn = sqlite3.connect(databasefile)
    cursor = conn.cursor()
    
    # SQL statements to create the tables
    create_links_table = """
    CREATE TABLE IF NOT EXISTS Links (
        id INTEGER PRIMARY KEY,
        source_ip TEXT NOT NULL,
        destination_ip TEXT NOT NULL,
        UNIQUE(source_ip, destination_ip)
    );
    """
    
    create_paths_table = """
    CREATE TABLE IF NOT EXISTS Paths (
        id INTEGER PRIMARY KEY,
        start_ip TEXT NOT NULL,
        end_ip TEXT NOT NULL,
        hops_json TEXT NOT NULL,
        UNIQUE(start_ip, end_ip, hops_json)
    );
    """
    
    create_latency_table = """
    CREATE TABLE IF NOT EXISTS Latency (
        id INTEGER PRIMARY KEY,
        link_id INTEGER NOT NULL,
        timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
        latency_ms REAL NOT NULL,
        FOREIGN KEY (link_id) REFERENCES Links(id)
    );
    """
    
    create_hopdetails_table = """
    CREATE TABLE IF NOT EXISTS HopDetails (
        id INTEGER PRIMARY KEY,
        link_id INTEGER NOT NULL,
        ttl INTEGER,
        icmp_type INTEGER,
        icmp_code INTEGER,
        packet_loss REAL,
        FOREIGN KEY (link_id) REFERENCES Links(id)
    );
    """
    
    # Execute the SQL statements
    cursor.execute(create_links_table)
    cursor.execute(create_paths_table)
    cursor.execute(create_latency_table)
    cursor.execute(create_hopdetails_table)
    
    # Commit changes and close the connection
    conn.commit()
    conn.close()

def store_traceroute(db_file, start_ip, end_ip, hops):
    conn = sqlite3.connect(db_file)
    cursor = conn.cursor()

    # Insert links and get their IDs
    link_ids = []
    for hop in hops:
        print(hop)
        source_ip = start_ip if not link_ids else hops[len(link_ids)-1][0]
        destination_ip = hop[0]
        latency = hop[1]

        cursor.execute("""
            INSERT OR IGNORE INTO Links (source_ip, destination_ip) VALUES (?, ?)
        """, (source_ip, destination_ip))
        cursor.execute("""
            SELECT id FROM Links WHERE source_ip = ? AND destination_ip = ?
        """, (source_ip, destination_ip))
        result = cursor.fetchone()
        if result is None:
            print(hop)
            raise ValueError(f"Failed to insert of find link between {source_ip} and {destination_ip}")
        link_id = result[0]
        link_ids.append(link_id)

        cursor.execute("""
            INSERT INTO Latency (link_id, latency_ms) VALUES (?, ?)
        """, (link_id, latency))

    # Insert path
    hops_json = json.dumps(link_ids)
    cursor.execute("""
        INSERT INTO Paths (start_ip, end_ip, hops_json) VALUES (?, ?, ?)
    """, (start_ip, end_ip, hops_json))

    conn.commit()
    conn.close()


def retrieve_traceroute():
    conn = sqlite3.connect('traceroute.db')
    cursor = conn.cursor()
    cursor.execute('''
        SELECT l.source_ip, l.destination_ip, ll.latency, ll.timestamp
        FROM link_latency ll
        JOIN links l ON ll.link_id = l.id
        ORDER BY ll.timestamp
    ''')
    rows = cursor.fetchall()
    conn.close()
    return rows

if __name__ == '__main__':

    databasefile="./traceroute.db"
    create_tables(databasefile)


    target='vi.fi'
    traceroute_output, timestamp = run_traceroute(target)
    hops = parse_traceroute_output(traceroute_output, timestamp)

    print("#####")
    print(hops)
    print("#####")

exit(0)