Update database schema for client tracking

- Add functions to manage client records (add, find, list)
- Require --confirm-add and --name flags when adding a new client
This commit is contained in:
Kris Lamoureux 2024-10-06 02:20:55 -04:00
parent 01e68bc10f
commit 236bb1b044
Signed by: kris
GPG Key ID: 3EDA9C3441EDA925

193
main.py
View File

@ -8,15 +8,20 @@ archival locations.
""" """
import argparse
import os import os
import sqlite3
import sys import sys
import re
import uuid
import argparse
import sqlite3
from datetime import datetime, timezone
import qbittorrent import qbittorrent
# SCHEMA format is YYYYMMDDX # SCHEMA format is YYYYMMDDX
SCHEMA = 202410040 SCHEMA = 202410060
def init_db(conn): def init_db(conn):
""" """
@ -27,36 +32,77 @@ def init_db(conn):
c.executescript( c.executescript(
f""" f"""
PRAGMA user_version = {SCHEMA}; PRAGMA user_version = {SCHEMA};
CREATE TABLE IF NOT EXISTS clients (
id INTEGER PRIMARY KEY AUTOINCREMENT,
name TEXT NOT NULL UNIQUE,
uuid TEXT NOT NULL UNIQUE,
endpoint TEXT NOT NULL,
last_seen DATETIME NOT NULL
);
CREATE TABLE IF NOT EXISTS torrents ( CREATE TABLE IF NOT EXISTS torrents (
id INTEGER PRIMARY KEY AUTOINCREMENT, id INTEGER PRIMARY KEY AUTOINCREMENT,
info_hash_v1 TEXT NOT NULL UNIQUE, info_hash_v1 TEXT NOT NULL UNIQUE,
info_hash_v2 TEXT UNIQUE, info_hash_v2 TEXT UNIQUE,
name TEXT NOT NULL,
file_count INTEGER NOT NULL, file_count INTEGER NOT NULL,
completed_on DATETIME NOT NULL
);
CREATE TABLE IF NOT EXISTS torrent_clients (
id INTEGER PRIMARY KEY AUTOINCREMENT,
torrent_id INTEGER NOT NULL,
client_id INTEGER NOT NULL,
name TEXT NOT NULL,
content_path TEXT NOT NULL, content_path TEXT NOT NULL,
completed_on DATETIME DEFAULT CURRENT_TIMESTAMP, last_seen DATETIME NOT NULL,
tracker_ids TEXT FOREIGN KEY (torrent_id) REFERENCES torrents(id),
FOREIGN KEY (client_id) REFERENCES clients(id),
UNIQUE (torrent_id, client_id)
); );
CREATE TABLE IF NOT EXISTS trackers ( CREATE TABLE IF NOT EXISTS trackers (
id INTEGER PRIMARY KEY AUTOINCREMENT, id INTEGER PRIMARY KEY AUTOINCREMENT,
url TEXT NOT NULL UNIQUE url TEXT NOT NULL UNIQUE,
last_seen DATETIME NOT NULL
);
CREATE TABLE IF NOT EXISTS torrent_trackers (
id INTEGER PRIMARY KEY AUTOINCREMENT,
client_id INTEGER NOT NULL,
torrent_id INTEGER NOT NULL,
tracker_id INTEGER NOT NULL,
last_seen DATETIME NOT NULL,
FOREIGN KEY (client_id) REFERENCES clients(id),
FOREIGN KEY (torrent_id) REFERENCES torrents(id),
FOREIGN KEY (tracker_id) REFERENCES trackers(id),
UNIQUE (client_id, torrent_id, tracker_id)
); );
CREATE TABLE IF NOT EXISTS files ( CREATE TABLE IF NOT EXISTS files (
id INTEGER PRIMARY KEY AUTOINCREMENT, id INTEGER PRIMARY KEY AUTOINCREMENT,
torrent_id INTEGER NOT NULL,
file_index INTEGER NOT NULL,
file_path TEXT NOT NULL,
size INTEGER NOT NULL, size INTEGER NOT NULL,
is_downloaded BOOLEAN NOT NULL DEFAULT 0, oshash TEXT NOT NULL UNIQUE,
last_checked DATETIME, hash TEXT UNIQUE
FOREIGN KEY (torrent_id) REFERENCES torrents(id),
UNIQUE (torrent_id, file_index)
); );
""" CREATE TABLE IF NOT EXISTS torrent_files (
id INTEGER PRIMARY KEY AUTOINCREMENT,
file_id INTEGER NOT NULL,
torrent_id INTEGER NOT NULL,
client_id INTEGER NOT NULL,
file_index INTEGER NOT NULL,
file_path TEXT NOT NULL,
is_downloaded BOOLEAN NOT NULL,
last_checked DATETIME NOT NULL,
FOREIGN KEY (file_id) REFERENCES files(id),
FOREIGN KEY (torrent_id) REFERENCES torrents(id),
FOREIGN KEY (client_id) REFERENCES clients(id),
UNIQUE (file_id, torrent_id, client_id, file_index)
);
"""
) )
conn.commit()
c.close() c.close()
@ -71,6 +117,43 @@ def list_tables(conn):
return [table[0] for table in table_list] return [table[0] for table in table_list]
def add_client(conn, name, endpoint, last_seen):
"""
Add a new client endpoint to database
"""
c = conn.cursor()
c.execute(
f"""
INSERT INTO clients (uuid, name, endpoint, last_seen)
VALUES ("{uuid.uuid4()}", "{name}", "{endpoint}", "{last_seen}");
"""
)
conn.commit()
c.close()
def find_client(conn, endpoint):
"""
Find existing client
"""
c = conn.cursor()
c.execute(f'SELECT id, name, uuid FROM clients WHERE endpoint="{endpoint}";')
response = c.fetchall()
c.close()
return response
def list_clients(conn):
"""
List all stored clients
"""
c = conn.cursor()
c.execute("SELECT * FROM clients;")
rows = c.fetchall()
c.close()
return rows
parser = argparse.ArgumentParser(description="Manage BitTorrent datasets", prog="tarch") parser = argparse.ArgumentParser(description="Manage BitTorrent datasets", prog="tarch")
subparsers = parser.add_subparsers( subparsers = parser.add_subparsers(
dest="command", required=True, help="Available commands" dest="command", required=True, help="Available commands"
@ -78,6 +161,10 @@ subparsers = parser.add_subparsers(
scan_parser = subparsers.add_parser("scan", help="Scan command") scan_parser = subparsers.add_parser("scan", help="Scan command")
scan_parser.add_argument("--debug", action="store_true", help="Enable debug mode") scan_parser.add_argument("--debug", action="store_true", help="Enable debug mode")
scan_parser.add_argument(
"--confirm-add", action="store_true", help="Confirm adding a new client"
)
scan_parser.add_argument("-n", "--name", help="Name of client")
scan_parser.add_argument("-d", "--directory", help="Directory to scan") scan_parser.add_argument("-d", "--directory", help="Directory to scan")
scan_parser.add_argument("-t", "--type", help="Scan type") scan_parser.add_argument("-t", "--type", help="Scan type")
scan_parser.add_argument("-e", "--endpoint", help="Endpoint URL") scan_parser.add_argument("-e", "--endpoint", help="Endpoint URL")
@ -96,7 +183,7 @@ if args.command == "scan":
sqlitedb = sqlite3.connect(STORAGE) sqlitedb = sqlite3.connect(STORAGE)
tables = list_tables(sqlitedb) tables = list_tables(sqlitedb)
except sqlite3.DatabaseError as e: except sqlite3.DatabaseError as e:
print(f"[ERROR]: Database \"{STORAGE}\" Error: {str(e)}") print(f'[ERROR]: Database Error "{STORAGE}" ({str(e)})')
sys.exit(1) sys.exit(1)
if len(tables) == 0: if len(tables) == 0:
print(f"[INFO]: Initializing database at {STORAGE}") print(f"[INFO]: Initializing database at {STORAGE}")
@ -109,19 +196,67 @@ if args.command == "scan":
print(f"[ERROR]: SCHEMA {SCHEMA_FOUND}, expected {SCHEMA}") print(f"[ERROR]: SCHEMA {SCHEMA_FOUND}, expected {SCHEMA}")
sys.exit(1) sys.exit(1)
if not args.directory is None: if not args.directory is None:
print("[INFO]: --directory is not implemented\n") print("[INFO]: --directory is not implemented")
sys.exit(0) sys.exit(0)
if not args.endpoint is None: elif not args.endpoint is None:
qb = qbittorrent.Client(args.endpoint) qb = qbittorrent.Client(args.endpoint)
torrents = qb.torrents() if qb.qbittorrent_version is None:
print(f"[INFO]: There are {len(torrents)} torrents\n") print(f'[ERROR]: Couldn\'t find client version at "{args.endpoint}"')
for torrent in torrents[:10]: sys.exit(1)
files = qb.get_torrent_files(torrent["hash"]) elif not re.match(r"^v?\d+(\.\d+)*$", qb.qbittorrent_version):
print(f'[ERROR]: Invalid version found at "{args.endpoint}"')
if args.debug: if args.debug:
print(f"[DEBUG]: {repr(torrent)}") print(f"[DEBUG]: {qb.qbittorrent_version}")
print(f"[name]: {torrent['name']}") sys.exit(1)
print(f"[infohash_v1]: {torrent['infohash_v1']}") else:
print(f"[content_path]: {torrent['content_path']}") print(
print(f"[magent_uri]: {torrent['magnet_uri'][0:80]}") f'[INFO]: Found qbittorrent {qb.qbittorrent_version} at "{args.endpoint}"'
print(f"[completed_on]: {torrent['completed']}") )
print(f"[file_count]: {len(files)}\n") clients = find_client(sqlitedb, args.endpoint)
if args.confirm_add:
if len(clients) == 0:
if not args.name is None:
now = datetime.now(timezone.utc).isoformat(
sep=" ", timespec="seconds"
)
add_client(sqlitedb, args.name, args.endpoint, now)
print(f"[INFO]: Added client {args.name} ({args.endpoint})")
else:
print("[ERROR]: Must specify --name for a new client")
sys.exit(1)
elif len(clients) == 1:
print(f"[ERROR]: {clients[0][1]} ({clients[0][2]}) already exists")
sys.exit(1)
else:
print(
f"[ERROR]: Multiple clients with the same endpoint: {args.endpoint}"
)
sys.exit(1)
elif len(clients) == 0:
print(f'[ERROR]: Client using endpoint "{args.endpoint}" not found')
print("[ERROR]: Use --confirm-add to add a new endpoint")
sys.exit(1)
elif len(clients) == 1:
torrents = qb.torrents()
print(f"[INFO]: There are {len(torrents)} torrents\n")
for torrent in torrents[:2]:
files = qb.get_torrent_files(torrent["hash"])
trackers = qb.get_torrent_trackers(torrent["hash"])
print(f"[name]: {torrent['name']}")
print(f"[infohash_v1]: {torrent['infohash_v1']}")
print(f"[content_path]: {torrent['content_path']}")
print(f"[magent_uri]: {torrent['magnet_uri'][0:80]}")
print(f"[completed_on]: {torrent['completed']}")
print(f"[trackers]: {len(trackers)}")
print(f"[file_count]: {len(files)}\n")
if args.debug:
print(f"[DEBUG]: {repr(torrent)}")
for elem in trackers:
print(f"[DEBUG]: Tracker {repr(elem)}")
print("\n", end="")
else:
print(f'[ERROR]: Multiple clients ({len(clients)}) using "{args.endpoint}"')
sys.exit(1)
else:
print("[ERROR]: Must specify directory OR client endpoint")
sys.exit(1)