traceroute_map/app/db.py

183 lines
5.0 KiB
Python

import sqlite3
from functools import wraps
# Type alias
Cursor = sqlite3.Cursor
# Configs
DB_FILE = "./traceroute.db"
class Database:
def __init__(self):
self.db_file = DB_FILE
self.conn = sqlite3.connect(self.db_file, check_same_thread=False)
self.cursor = self.conn.cursor()
def create_tables(self):
self.cursor.executescript("""
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 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 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 TABLE IF NOT EXISTS HopDetails (
id INTEGER PRIMARY KEY,
hop_name TEXT,
hop_ip TEXT,
hop_latency TEXT
);
""")
def end(self):
"""Always call this after you're done with the connection / request."""
self.conn.commit()
self.conn.close()
def get_traceroute(self):
retval = {}
self.cursor.execute('''
SELECT target, hops_json
FROM Paths
''')
retval['path'] = self.cursor.fetchall()
self.cursor.execute('''
SELECT source_ip, destination_ip
FROM Links
''')
retval['links'] = self.cursor.fetchall()
return retval
def create_link(self, previous_hop_ip, hop_ip):
self.cursor.execute(
"INSERT OR IGNORE INTO Links (source_ip, destination_ip) VALUES (?, ?)",
(previous_hop_ip, hop_ip)
)
self.cursor.execute(
"SELECT id FROM Links WHERE source_ip = ? AND destination_ip = ?",
(previous_hop_ip, hop_ip)
)
return self.cursor.fetchone()
def create_hop(self, name, ip, latency):
self.cursor.execute(
"INSERT INTO HopDetails (hop_name, hop_ip, hop_latency) VALUES (?, ?, ?)",
(name, ip, latency)
)
def create_latency(self, link_id, timestamp, link_latency):
self.cursor.execute(
"INSERT INTO Latency (link_id, timestamp, latency_ms) VALUES (?, NOW(), ?)",
(link_id, timestamp, link_latency)
)
def create_path(self, node, target, json):
self.cursor.execute(
"INSERT OR IGNORE INTO Paths (node, target, hops_json) VALUES (?, ?, ?)",
(node, target, json)
)
def ensure_table_setup():
db = Database()
db.create_tables()
db.end()
####################################################################
####################################################################
####################################################################
####################################################################
# Temp testing. Fancy decorator stuff.
def with_connection(func):
@wraps(func)
def wrapped(*args, **kwargs):
conn = sqlite3.connect(DB_FILE)
cursor = conn.cursor()
result = func(cursor, *args, **kwargs)
conn.commit()
conn.close()
return result
return wrapped
@with_connection
def init_db(cursor: Cursor):
cursor.executescript("""
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 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 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 TABLE IF NOT EXISTS HopDetails (
id INTEGER PRIMARY KEY,
hop_name TEXT,
hop_ip TEXT,
hop_latency TEXT
);
""")
@with_connection
def create_link(cursor: Cursor, previous_hop_ip: str, hop_ip: str):
"""Insert a new hop and return related Link id"""
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))
return cursor.fetchone()