#!/usr/bin/env python3.11


import subprocess
import sqlite3
import re
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):
    hops = []
    lines = output.strip().split('\n')[1:]  # Skip the first line (traceroute to ...)
    previous_ip = None
    for line in lines:
        parts = line.split()
        hop_number = int(parts[0])
        ip_address = parts[1]
        latencies = []
        for part in parts[2:]:
            if re.match(r'^[0-9+\s+ms$]', part):
                latency_str = re.sub(r'[^0-9.]', '', part)
                if latency_str and not latency_str == '':
                    try:
                        print(part)
                        print(latency_str)
                        latencies.append(float(latency_str))
                    except ValueError:
                        print(f"Could not convert '{latency_str}' to float.")

        avg_latency = sum(latencies) / len(latencies) if latencies else None
        timestamp = datetime.now().isoformat()
        if previous_ip:
            hops.append({
                'hop_number': hop_number,
                'source_ip': previous_ip,
                'destination_ip': ip_address,
                'latency': avg_latency,
                'timestamp': timestamp,
            })
        previous_ip = ip_address
    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(hops):
    conn = sqlite3.connect('traceroute.db')
    cursor = conn.cursor()
    
    for hop in hops:
        # Insert or ignore link into links table
        cursor.execute('''
            INSERT OR IGNORE INTO links (source_ip, destination_ip)
            VALUES (?, ?)
        ''', (hop['source_ip'], hop['destination_ip']))
        
        # Retrieve the link_id
        cursor.execute('''
            SELECT id FROM links WHERE source_ip = ? AND destination_ip = ?
        ''', (hop['source_ip'], hop['destination_ip']))
        link_id = cursor.fetchone()[0]
        
        # Insert latency data into link_latency table
        cursor.execute('''
            INSERT INTO link_latency (link_id, latency, timestamp)
            VALUES (?, ?, ?)
        ''', (link_id, hop['latency'], hop['timestamp']))
    
    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

# Usage
if __name__ == '__main__':
    create_tables()
    traceroute_output = run_traceroute('vi.fi')
    hops = parse_traceroute_output(traceroute_output)
#    for hop in hops:
#        print(hop)
    store_traceroute(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)