forked from MrKalzu/traceroute_map
118 lines
3.1 KiB
Python
118 lines
3.1 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)
|
|
|
|
# Return fetch() data as Row objects, instead of tuples.
|
|
self.conn.row_factory = sqlite3.Row
|
|
|
|
self.cursor = self.conn.cursor()
|
|
|
|
def create_tables(self):
|
|
self.cursor.executescript(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS Traces (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
created TEXT NOT NULL,
|
|
origin TEXT NOT NULL,
|
|
target TEXT NOT NULL,
|
|
unparsed TEXT NOT NULL
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS Nodes (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
created TEXT NOT NULL,
|
|
name TEXT,
|
|
ip TEXT,
|
|
latency_ms REAL NOT NULL,
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS Links (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
created TEXT NOT NULL,
|
|
source TEXT NOT NULL,
|
|
target TEXT NOT NULL,
|
|
latency_ms REAL NOT NULL,
|
|
);
|
|
"""
|
|
)
|
|
|
|
def end(self):
|
|
"""Always call this after you're done with the connection / request."""
|
|
|
|
self.conn.commit()
|
|
self.conn.close()
|
|
|
|
# TODO
|
|
def list_traces(self):
|
|
# TODO: time filter
|
|
result = []
|
|
|
|
self.cursor.execute("SELECT * FROM Traces")
|
|
traces = self.cursor.fetchall()
|
|
|
|
for t in traces:
|
|
trace = dict(t)
|
|
|
|
self.cursor.execute(
|
|
"""
|
|
SELECT number, name, ip, latency, link_latency
|
|
FROM Hops
|
|
WHERE trace_id = ?
|
|
ORDER BY number ASC
|
|
""",
|
|
(trace["id"],),
|
|
)
|
|
hops = self.cursor.fetchall()
|
|
trace["hops"] = hops
|
|
|
|
result.append(trace)
|
|
|
|
return result
|
|
|
|
# TODO
|
|
def create_trace(self, trace):
|
|
self.cursor.execute(
|
|
"INSERT OR IGNORE INTO Traces (created, origin, target) VALUES (?, ?, ?)",
|
|
(trace["created"], trace["origin"], trace["target"]),
|
|
)
|
|
trace_id = self.cursor.lastrowid
|
|
|
|
for hop in trace["hops"]:
|
|
self.cursor.execute(
|
|
"INSERT OR IGNORE INTO Hops (trace_id, created, number, name, ip, latency, link_latency) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
|
(
|
|
trace_id,
|
|
hop["created"], # TODO: trace.created
|
|
hop["number"],
|
|
hop["name"],
|
|
hop["ip"],
|
|
hop["latency"],
|
|
hop["link_latency"],
|
|
),
|
|
)
|
|
|
|
return self.cursor.fetchone()
|
|
|
|
def create_hop(self, name, ip, latency):
|
|
self.cursor.execute(
|
|
"INSERT INTO Hops (name, ip, latency) VALUES (?, ?, ?)",
|
|
(name, ip, latency),
|
|
)
|
|
|
|
|
|
def ensure_table_setup():
|
|
db = Database()
|
|
db.create_tables()
|
|
db.end()
|