#!/usr/bin/env python3.11


import subprocess
import sqlite3
import re
import json
import ipaddress
import uuid
import hashlib
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')
    trace = {}
    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 = 'unresponsive'
                hop_ip = 'unresponsive'
                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['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(hop)
        
    trace['target'] = target
    trace['timestamp'] = timestamp
    trace['hops'] = hops

    return trace


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,
        node TEXT NOT NULL,
        target TEXT NOT NULL,
        hops_json TEXT NOT NULL,
        UNIQUE(node, target, 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,
        hop_name TEXT,
        hop_ip TEXT,
        hop_latency TEXT
    );
    """
    
    # 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, node, trace):
    conn = sqlite3.connect(db_file)
    cursor = conn.cursor()


    hops_json = json.dumps(trace['hops'])
    
    path_ids = {}

    previous_hop_ip = None
    previous_hop_latency = None
    for hop in trace['hops']:
        hop_number = hop['hop_number']
        hop_name = hop.get('hop_name')
        hop_ip = hop.get('hop_ip')
        hop_latency = hop.get('hop_latency')
        link_id = None

        # insert links and get their id's
        if previous_hop_ip:
            cursor.execute("""
                INSERT OR IGNORE INTO Links (source_ip, destination_ip) VALUES (?, ?)
            """, (previous_hop_ip, hop_ip))
            cursor.execute("""
                SELECT id FROM Links WHERE source_ip = ? AND destination_ip = ?
            """, (previous_hop_ip, hop_ip))
            result = cursor.fetchone()
            link_id = result
            path_ids[hop_number] = link_id
        previous_hop_ip = hop_ip

        # Save hop details
        cursor.execute("INSERT INTO HopDetails (hop_name, hop_ip, hop_latency) VALUES (?, ?, ?)",
                        (hop_name, hop_ip, hop_latency))
        
        # calculate link latency if possible and store it
        if link_id and previous_hop_latency:
            link_latency = hop_latency - previous_hop_latency
            cursor.execute("INSERT INTO Latency (link_id, timestamp, latency_ms) VALUES (?, ?, ?)", 
                           (link_id, trace['timestamp'], link_latency))

    # make entry to "Paths" table
    if path_ids:
        json_path_ids = json.dumps(path_ids)
        cursor.execute("INSERT OR IGNORE INTO Paths (node, target, hops_json) VALUES (?, ?, ?)",
                       (node, trace['target'], json_path_ids))

    conn.commit()
    conn.close()


def retrieve_traceroute(db_file):

    retval = {}

    conn = sqlite3.connect(db_file)
    cursor = conn.cursor()
    cursor.execute('''
        SELECT target, hops_json
        FROM Paths
    ''')
    retval['path'] = cursor.fetchall()

    cursor.execute('''
        SELECT source_ip, destination_ip
        FROM Links
    ''')
    retval['links'] = cursor.fetchall()
    conn.close()
    return retval

def generate_node_id():
    mac = uuid.getnode()
    mac_str = ':'.join(['{:02x}'.format((mac >> ele) & 0xff) for ele in range(0,8*6,8)][::-1])
    
    # Hash the MAC address using SHA-256 to generate a unique ID
    unique_id = hashlib.sha256(mac_str.encode()).hexdigest()

    return unique_id


if __name__ == '__main__':

    db_file="./traceroute.db"
    create_tables(db_file)
    my_id = generate_node_id()

    target='1.1.1.1'
    traceroute_output, timestamp = run_traceroute(target)
    trace = parse_traceroute_output(traceroute_output, timestamp)

    store_traceroute(db_file, my_id, trace)

    print("#####")
    print(retrieve_traceroute(db_file))
    print("#####")

exit(0)