From 236bb1b044868a9da816e478c91ac4df2c09a926 Mon Sep 17 00:00:00 2001 From: Kris Lamoureux Date: Sun, 6 Oct 2024 02:20:55 -0400 Subject: [PATCH] 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 --- main.py | 193 +++++++++++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 164 insertions(+), 29 deletions(-) diff --git a/main.py b/main.py index ff2646c..6fa0a84 100644 --- a/main.py +++ b/main.py @@ -8,15 +8,20 @@ archival locations. """ -import argparse import os -import sqlite3 import sys +import re +import uuid +import argparse +import sqlite3 + +from datetime import datetime, timezone import qbittorrent # SCHEMA format is YYYYMMDDX -SCHEMA = 202410040 +SCHEMA = 202410060 + def init_db(conn): """ @@ -27,36 +32,77 @@ def init_db(conn): c.executescript( f""" 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 ( id INTEGER PRIMARY KEY AUTOINCREMENT, info_hash_v1 TEXT NOT NULL UNIQUE, info_hash_v2 TEXT UNIQUE, - name TEXT 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, - completed_on DATETIME DEFAULT CURRENT_TIMESTAMP, - tracker_ids TEXT + last_seen DATETIME NOT NULL, + 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 ( 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 ( id INTEGER PRIMARY KEY AUTOINCREMENT, - torrent_id INTEGER NOT NULL, - file_index INTEGER NOT NULL, - file_path TEXT NOT NULL, size INTEGER NOT NULL, - is_downloaded BOOLEAN NOT NULL DEFAULT 0, - last_checked DATETIME, - FOREIGN KEY (torrent_id) REFERENCES torrents(id), - UNIQUE (torrent_id, file_index) + oshash TEXT NOT NULL UNIQUE, + hash TEXT UNIQUE ); - """ + 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() @@ -71,6 +117,43 @@ def list_tables(conn): 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") subparsers = parser.add_subparsers( 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.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("-t", "--type", help="Scan type") scan_parser.add_argument("-e", "--endpoint", help="Endpoint URL") @@ -96,7 +183,7 @@ if args.command == "scan": sqlitedb = sqlite3.connect(STORAGE) tables = list_tables(sqlitedb) except sqlite3.DatabaseError as e: - print(f"[ERROR]: Database \"{STORAGE}\" Error: {str(e)}") + print(f'[ERROR]: Database Error "{STORAGE}" ({str(e)})') sys.exit(1) if len(tables) == 0: print(f"[INFO]: Initializing database at {STORAGE}") @@ -109,19 +196,67 @@ if args.command == "scan": print(f"[ERROR]: SCHEMA {SCHEMA_FOUND}, expected {SCHEMA}") sys.exit(1) if not args.directory is None: - print("[INFO]: --directory is not implemented\n") + print("[INFO]: --directory is not implemented") sys.exit(0) - if not args.endpoint is None: + elif not args.endpoint is None: qb = qbittorrent.Client(args.endpoint) - torrents = qb.torrents() - print(f"[INFO]: There are {len(torrents)} torrents\n") - for torrent in torrents[:10]: - files = qb.get_torrent_files(torrent["hash"]) + if qb.qbittorrent_version is None: + print(f'[ERROR]: Couldn\'t find client version at "{args.endpoint}"') + sys.exit(1) + elif not re.match(r"^v?\d+(\.\d+)*$", qb.qbittorrent_version): + print(f'[ERROR]: Invalid version found at "{args.endpoint}"') if args.debug: - print(f"[DEBUG]: {repr(torrent)}") - 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"[file_count]: {len(files)}\n") + print(f"[DEBUG]: {qb.qbittorrent_version}") + sys.exit(1) + else: + print( + f'[INFO]: Found qbittorrent {qb.qbittorrent_version} at "{args.endpoint}"' + ) + 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)