From 7b69d95b54dd7657cfe99e2fc17127fe2cb1a0ef Mon Sep 17 00:00:00 2001
From: Kris Lamoureux <kris@lamoureux.io>
Date: Sun, 6 Oct 2024 01:43:33 -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 | 181 ++++++++++++++++++++++++++++++++++++++++++++++----------
1 file changed, 151 insertions(+), 30 deletions(-)
diff --git a/main.py b/main.py
index ff2646c..9021258 100644
--- a/main.py
+++ b/main.py
@@ -8,15 +8,19 @@ archival locations.
"""
-import argparse
import os
-import sqlite3
import sys
+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 +31,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 +116,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 +160,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 +182,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 +195,54 @@ 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 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")
+ 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"[tracker]: {repr(elem)}")
+ 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)