From 5637ff3b197d732d5a4bfe76d0153aee515e88c7 Mon Sep 17 00:00:00 2001 From: Kris Lamoureux Date: Fri, 28 Feb 2025 21:00:50 -0500 Subject: [PATCH] testing --- requirements.txt | 1 + tarc/main.py | 266 ++++++++++++++++++++++++++--------------------- 2 files changed, 147 insertions(+), 120 deletions(-) diff --git a/requirements.txt b/requirements.txt index 1ab3894..3316d12 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ python-qbittorrent==0.4.3 +sqlalchemy==2.0.27 diff --git a/tarc/main.py b/tarc/main.py index 1432e5e..4e5fbe0 100644 --- a/tarc/main.py +++ b/tarc/main.py @@ -13,145 +13,167 @@ import sys import re import uuid import argparse -import sqlite3 - from datetime import datetime, timezone import qbittorrent +from sqlalchemy import create_engine, Column, Integer, String, Boolean, DateTime, ForeignKey, UniqueConstraint, text +from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy.orm import sessionmaker, relationship # SCHEMA format is YYYYMMDDX SCHEMA = 202410060 +Base = declarative_base() -def init_db(conn): + +class Client(Base): + __tablename__ = 'clients' + + id = Column(Integer, primary_key=True, autoincrement=True) + name = Column(String, nullable=False, unique=True) + uuid = Column(String, nullable=False, unique=True) + endpoint = Column(String, nullable=False) + last_seen = Column(DateTime, nullable=False) + + torrent_clients = relationship("TorrentClient", back_populates="client") + torrent_trackers = relationship("TorrentTracker", back_populates="client") + torrent_files = relationship("TorrentFile", back_populates="client") + + +class Torrent(Base): + __tablename__ = 'torrents' + + id = Column(Integer, primary_key=True, autoincrement=True) + info_hash_v1 = Column(String, nullable=False, unique=True) + info_hash_v2 = Column(String, unique=True) + file_count = Column(Integer, nullable=False) + completed_on = Column(DateTime, nullable=False) + + torrent_clients = relationship("TorrentClient", back_populates="torrent") + torrent_trackers = relationship("TorrentTracker", back_populates="torrent") + torrent_files = relationship("TorrentFile", back_populates="torrent") + + +class TorrentClient(Base): + __tablename__ = 'torrent_clients' + + id = Column(Integer, primary_key=True, autoincrement=True) + torrent_id = Column(Integer, ForeignKey('torrents.id'), nullable=False) + client_id = Column(Integer, ForeignKey('clients.id'), nullable=False) + name = Column(String, nullable=False) + content_path = Column(String, nullable=False) + last_seen = Column(DateTime, nullable=False) + + torrent = relationship("Torrent", back_populates="torrent_clients") + client = relationship("Client", back_populates="torrent_clients") + + __table_args__ = (UniqueConstraint('torrent_id', 'client_id'),) + + +class Tracker(Base): + __tablename__ = 'trackers' + + id = Column(Integer, primary_key=True, autoincrement=True) + url = Column(String, nullable=False, unique=True) + last_seen = Column(DateTime, nullable=False) + + torrent_trackers = relationship("TorrentTracker", back_populates="tracker") + + +class TorrentTracker(Base): + __tablename__ = 'torrent_trackers' + + id = Column(Integer, primary_key=True, autoincrement=True) + client_id = Column(Integer, ForeignKey('clients.id'), nullable=False) + torrent_id = Column(Integer, ForeignKey('torrents.id'), nullable=False) + tracker_id = Column(Integer, ForeignKey('trackers.id'), nullable=False) + last_seen = Column(DateTime, nullable=False) + + client = relationship("Client", back_populates="torrent_trackers") + torrent = relationship("Torrent", back_populates="torrent_trackers") + tracker = relationship("Tracker", back_populates="torrent_trackers") + + __table_args__ = (UniqueConstraint('client_id', 'torrent_id', 'tracker_id'),) + + +class File(Base): + __tablename__ = 'files' + + id = Column(Integer, primary_key=True, autoincrement=True) + size = Column(Integer, nullable=False) + oshash = Column(String, nullable=False, unique=True) + hash = Column(String, unique=True) + + torrent_files = relationship("TorrentFile", back_populates="file") + + +class TorrentFile(Base): + __tablename__ = 'torrent_files' + + id = Column(Integer, primary_key=True, autoincrement=True) + file_id = Column(Integer, ForeignKey('files.id'), nullable=False) + torrent_id = Column(Integer, ForeignKey('torrents.id'), nullable=False) + client_id = Column(Integer, ForeignKey('clients.id'), nullable=False) + file_index = Column(Integer, nullable=False) + file_path = Column(String, nullable=False) + is_downloaded = Column(Boolean, nullable=False) + last_checked = Column(DateTime, nullable=False) + + file = relationship("File", back_populates="torrent_files") + torrent = relationship("Torrent", back_populates="torrent_files") + client = relationship("Client", back_populates="torrent_files") + + __table_args__ = (UniqueConstraint('file_id', 'torrent_id', 'client_id', 'file_index'),) + + +def init_db(engine): """ Initialize database """ - - c = conn.cursor() - 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, - 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, - 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, - 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, - size INTEGER NOT NULL, - 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() + # Set the schema version + with engine.connect() as conn: + conn.execute(text(f"PRAGMA user_version = {SCHEMA}")) + + # Create all tables + Base.metadata.create_all(engine) -def list_tables(conn): +def list_tables(engine): """ List all tables in database """ - c = conn.cursor() - c.execute("SELECT name FROM sqlite_master WHERE type='table';") - table_list = c.fetchall() - c.close() - return [table[0] for table in table_list] + with engine.connect() as conn: + result = conn.execute(text("SELECT name FROM sqlite_master WHERE type='table';")) + return [row[0] for row in result] -def add_client(conn, name, endpoint, last_seen): +def add_client(session, 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}"); - """ + new_client = Client( + uuid=str(uuid.uuid4()), + name=name, + endpoint=endpoint, + last_seen=last_seen ) - conn.commit() - c.close() + session.add(new_client) + session.commit() -def find_client(conn, endpoint): +def find_client(session, 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 + clients = session.query(Client.id, Client.name, Client.uuid).filter(Client.endpoint == endpoint).all() + return clients -def list_clients(conn): +def list_clients(session): """ List all stored clients """ - c = conn.cursor() - c.execute("SELECT * FROM clients;") - rows = c.fetchall() - c.close() - return rows + return session.query(Client).all() def main(): @@ -185,18 +207,21 @@ def main(): else: STORAGE = args.storage try: - sqlitedb = sqlite3.connect(STORAGE) - tables = list_tables(sqlitedb) - except sqlite3.DatabaseError as e: + engine = create_engine(f"sqlite:///{STORAGE}") + tables = list_tables(engine) + Session = sessionmaker(bind=engine) + session = Session() + except Exception as e: print(f'[ERROR]: Database Error "{STORAGE}" ({str(e)})') sys.exit(1) if len(tables) == 0: print(f"[INFO]: Initializing database at {STORAGE}") - init_db(sqlitedb) - cursor = sqlitedb.cursor() - cursor.execute("PRAGMA user_version;") - SCHEMA_FOUND = cursor.fetchone()[0] - cursor.close() + init_db(engine) + + # Check schema version + with engine.connect() as conn: + SCHEMA_FOUND = conn.execute(text("PRAGMA user_version;")).fetchone()[0] + if not SCHEMA == SCHEMA_FOUND: print(f"[ERROR]: SCHEMA {SCHEMA_FOUND}, expected {SCHEMA}") sys.exit(1) @@ -205,6 +230,8 @@ def main(): sys.exit(0) elif not args.endpoint is None: qb = qbittorrent.Client(args.endpoint) + if args.username and args.password: + qb.login(args.username, args.password) if qb.qbittorrent_version is None: print(f'[ERROR]: Couldn\'t find client version at "{args.endpoint}"') sys.exit(1) @@ -217,14 +244,12 @@ def main(): print( f'[INFO]: Found qbittorrent {qb.qbittorrent_version} at "{args.endpoint}"' ) - clients = find_client(sqlitedb, args.endpoint) + clients = find_client(session, 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) + now = datetime.now(timezone.utc) + add_client(session, args.name, args.endpoint, now) print(f"[INFO]: Added client {args.name} ({args.endpoint})") else: print("[ERROR]: Must specify --name for a new client") @@ -244,6 +269,7 @@ def main(): 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"])