#!/usr/bin/env python3.11


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

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

def parse_traceroute_output(output):

    lines = output.strip().split('\n')
    hops = []
    for line in lines[1:]:
        hop_info = line.split()
        hop_number = int(hop_info[0])
        hop_ips = []
        hop_latencies = []
        for part in hop_info[1:]:
            if re.match(r'\d+\.\d+\.\d+\.\d+', part):  # Match IP address
                hop_ips.append(part)
            elif re.match(r'\d+(\.\d+)? ms', part):  # Match latency
                hop_latencies.append(float(part.replace(' ms', '')))
        
        # If multiple IPs are present, we consider the first as primary
        primary_ip = hop_ips[0] if hop_ips else None
        avg_latency = sum(hop_latencies) / len(hop_latencies) if hop_latencies else None
        hops.append((primary_ip, avg_latency))

    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:
        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))
        link_id = cursor.fetchone()[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 = run_traceroute(target)
    hops = parse_traceroute_output(traceroute_output)
    if hops:
        start_ip = hops[0][0]
        store_traceroute(databasefile, start_ip, target, hops)
#    stored_hops = retrieve_traceroute()
#    for hop in stored_hops:
#        print(f"Link: {hop[0]} -> {hop[1]}, Latency: {hop[2]} ms, Timestamp: {hop[3]}")


exit(0)