bvandeusen
5bcb35e737
Files within a downloading torrent are now imported as soon as their individual progress reaches 100%, without waiting for the full torrent to complete. A new processed_files table (info_hash, file_name) tracks which files have been handled to avoid redundant re-hashing each poll. - db.py: add processed_files table, is_file_processed, record_processed_file - qbit_client.py: replace get_completed_torrents with get_torrents(filter, category) - stash_handler.py: extract process_torrent_files helper, daemon now polls both downloading and completed torrents each cycle Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
278 lines
8.8 KiB
Python
278 lines
8.8 KiB
Python
#!/usr/bin/env python3
|
|
"""StashHandler — bridge between qBittorrent and Stash."""
|
|
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import signal
|
|
import shutil
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import db
|
|
from hasher import hash_file
|
|
from qbit_client import QBitClient
|
|
|
|
log = logging.getLogger("stash_handler")
|
|
|
|
DEFAULT_MEDIA_EXTENSIONS = ".mp4,.mkv,.avi,.wmv,.mov,.webm,.flv,.m4v"
|
|
|
|
|
|
def load_config():
|
|
"""Build config dict from environment variables."""
|
|
return {
|
|
"qbit_url": os.environ.get("QBIT_URL", "http://qbittorrent:8080"),
|
|
"qbit_username": os.environ.get("QBIT_USERNAME", "admin"),
|
|
"qbit_password": os.environ.get("QBIT_PASSWORD", "adminadmin"),
|
|
"staging_dir": os.environ.get("STAGING_DIR", "/staging"),
|
|
"import_dir": os.environ.get("IMPORT_DIR", "/import"),
|
|
"poll_interval": int(os.environ.get("POLL_INTERVAL", "60")),
|
|
"media_extensions": set(
|
|
os.environ.get("MEDIA_EXTENSIONS", DEFAULT_MEDIA_EXTENSIONS).split(",")
|
|
),
|
|
"hash_algorithm": os.environ.get("HASH_ALGORITHM", "sha256"),
|
|
"move_mode": os.environ.get("MOVE_MODE", "copy"),
|
|
"qbit_category": os.environ.get("QBIT_CATEGORY", "Stash"),
|
|
"db_path": os.environ.get("DB_PATH", "/data/seen.db"),
|
|
"log_level": os.environ.get("LOG_LEVEL", "INFO"),
|
|
}
|
|
|
|
|
|
def setup_logging(level="INFO"):
|
|
logging.basicConfig(
|
|
level=getattr(logging, level.upper(), logging.INFO),
|
|
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
)
|
|
|
|
|
|
def is_media(path, extensions):
|
|
return path.suffix.lower() in extensions
|
|
|
|
|
|
def walk_media(root, extensions):
|
|
"""Yield all media file paths under root."""
|
|
root = Path(root)
|
|
for p in sorted(root.rglob("*")):
|
|
if p.is_file() and is_media(p, extensions):
|
|
yield p
|
|
|
|
|
|
def transfer_file(src, dst, mode):
|
|
"""Move, hardlink, or copy src to dst."""
|
|
dst.parent.mkdir(parents=True, exist_ok=True)
|
|
if mode == "hardlink":
|
|
try:
|
|
os.link(src, dst)
|
|
except OSError:
|
|
log.warning("Hardlink failed, falling back to copy: %s", src)
|
|
shutil.copy2(src, dst)
|
|
elif mode == "copy":
|
|
shutil.copy2(src, dst)
|
|
else: # move
|
|
shutil.move(str(src), str(dst))
|
|
|
|
|
|
# --- Daemon mode ---
|
|
|
|
def process_torrent_files(conn, qbit, t, staging, import_dir, extensions, algorithm, move_mode):
|
|
"""
|
|
Process all complete files for a torrent using the qBit files API.
|
|
|
|
Files with progress < 1.0 are skipped (still downloading).
|
|
Files already in processed_files are skipped (already handled in a prior poll).
|
|
Returns (new_count, dup_count, incomplete_count).
|
|
"""
|
|
info_hash = t["hash"]
|
|
name = t["name"]
|
|
new_count = 0
|
|
dup_count = 0
|
|
incomplete_count = 0
|
|
|
|
try:
|
|
files = qbit.get_torrent_files(info_hash)
|
|
except Exception:
|
|
log.exception("Failed to get file list for torrent: %s", name)
|
|
return 0, 0, 0
|
|
|
|
for f in files:
|
|
file_name = f["name"] # relative to qBit's save_path (= our staging root)
|
|
|
|
if f["progress"] < 1.0:
|
|
incomplete_count += 1
|
|
continue
|
|
|
|
if db.is_file_processed(conn, info_hash, file_name):
|
|
continue
|
|
|
|
fpath = staging / file_name
|
|
if not fpath.exists():
|
|
log.warning("File reported complete by qBit but not found on disk: %s", fpath)
|
|
continue
|
|
|
|
if not is_media(fpath, extensions):
|
|
# Not a media file — mark processed so we don't check it again
|
|
db.record_processed_file(conn, info_hash, file_name)
|
|
continue
|
|
|
|
file_hash = hash_file(fpath, algorithm)
|
|
if db.is_known(conn, file_hash):
|
|
db.increment_seen(conn, file_hash)
|
|
dup_count += 1
|
|
log.debug("Duplicate skipped: %s", fpath.name)
|
|
else:
|
|
rel = Path(file_name)
|
|
dst = import_dir / rel
|
|
transfer_file(fpath, dst, move_mode)
|
|
db.record_file(conn, file_hash, str(rel), "torrent", info_hash)
|
|
new_count += 1
|
|
log.info("Imported: %s", rel)
|
|
|
|
db.record_processed_file(conn, info_hash, file_name)
|
|
|
|
return new_count, dup_count, incomplete_count
|
|
|
|
|
|
def daemon(cfg):
|
|
qbit = QBitClient(cfg["qbit_url"], cfg["qbit_username"], cfg["qbit_password"])
|
|
conn = db.init_db(cfg["db_path"])
|
|
extensions = cfg["media_extensions"]
|
|
staging = Path(cfg["staging_dir"])
|
|
import_dir = Path(cfg["import_dir"])
|
|
algorithm = cfg["hash_algorithm"]
|
|
move_mode = cfg["move_mode"]
|
|
poll_interval = cfg["poll_interval"]
|
|
category = cfg["qbit_category"]
|
|
|
|
running = True
|
|
|
|
def stop(*_):
|
|
nonlocal running
|
|
log.info("Shutting down...")
|
|
running = False
|
|
|
|
signal.signal(signal.SIGINT, stop)
|
|
signal.signal(signal.SIGTERM, stop)
|
|
|
|
log.info("Daemon started — polling every %ds (category: %s)", poll_interval, category)
|
|
|
|
while running:
|
|
try:
|
|
downloading = qbit.get_torrents(filter="downloading", category=category)
|
|
completed = qbit.get_torrents(filter="completed", category=category)
|
|
except Exception:
|
|
log.exception("Failed to fetch torrents from qBittorrent")
|
|
_sleep(poll_interval, lambda: running)
|
|
continue
|
|
|
|
# Process in-progress torrents — import any individually complete files
|
|
for t in downloading:
|
|
if not running:
|
|
break
|
|
if db.is_torrent_processed(conn, t["hash"]):
|
|
continue
|
|
new, dup, incomplete = process_torrent_files(
|
|
conn, qbit, t, staging, import_dir, extensions, algorithm, move_mode
|
|
)
|
|
if new or dup:
|
|
log.info(
|
|
"In-progress '%s': %d imported, %d duplicates, %d still downloading",
|
|
t["name"], new, dup, incomplete,
|
|
)
|
|
|
|
# Process completed torrents — import remaining files and mark torrent done
|
|
for t in completed:
|
|
if not running:
|
|
break
|
|
if db.is_torrent_processed(conn, t["hash"]):
|
|
continue
|
|
log.info("Completed torrent: %s (%s)", t["name"], t["hash"])
|
|
new, dup, _ = process_torrent_files(
|
|
conn, qbit, t, staging, import_dir, extensions, algorithm, move_mode
|
|
)
|
|
db.record_torrent(conn, t["hash"], t["name"])
|
|
log.info("Torrent done: '%s' — %d imported, %d duplicates", t["name"], new, dup)
|
|
|
|
_sleep(poll_interval, lambda: running)
|
|
|
|
conn.close()
|
|
log.info("Shutdown complete")
|
|
|
|
|
|
def _sleep(seconds, check_fn, step=1):
|
|
"""Sleep in small steps so we can respond to shutdown quickly."""
|
|
elapsed = 0
|
|
while elapsed < seconds and check_fn():
|
|
time.sleep(min(step, seconds - elapsed))
|
|
elapsed += step
|
|
|
|
|
|
# --- Scan mode ---
|
|
|
|
def scan(cfg, path, do_import):
|
|
conn = db.init_db(cfg["db_path"])
|
|
extensions = cfg["media_extensions"]
|
|
import_dir = Path(cfg["import_dir"])
|
|
algorithm = cfg["hash_algorithm"]
|
|
move_mode = cfg["move_mode"]
|
|
scan_root = Path(path)
|
|
|
|
total = 0
|
|
new = 0
|
|
duplicates = 0
|
|
|
|
for fpath in walk_media(scan_root, extensions):
|
|
total += 1
|
|
file_hash = hash_file(fpath, algorithm)
|
|
|
|
if db.is_known(conn, file_hash):
|
|
db.increment_seen(conn, file_hash)
|
|
duplicates += 1
|
|
log.debug("Already known: %s", fpath.name)
|
|
else:
|
|
if do_import:
|
|
rel = fpath.relative_to(scan_root)
|
|
dst = import_dir / rel
|
|
transfer_file(fpath, dst, move_mode)
|
|
log.info("Imported: %s", rel)
|
|
db.record_file(conn, file_hash, str(fpath), "scan" if do_import else "bulk")
|
|
new += 1
|
|
|
|
conn.close()
|
|
|
|
print(f"\nScan complete: {path}")
|
|
print(f" Total files: {total}")
|
|
print(f" New: {new}")
|
|
print(f" Already known: {duplicates}")
|
|
if do_import:
|
|
print(f" Imported to: {import_dir}")
|
|
else:
|
|
print(" (record-only mode — no files moved)")
|
|
|
|
|
|
# --- CLI ---
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="StashHandler")
|
|
sub = parser.add_subparsers(dest="command")
|
|
|
|
sub.add_parser("daemon", help="Poll qBittorrent and import new files")
|
|
|
|
scan_p = sub.add_parser("scan", help="Bulk-scan files and record/import them")
|
|
scan_p.add_argument("path", nargs="?", default="/staging", help="Directory to scan")
|
|
scan_p.add_argument("--import", dest="do_import", action="store_true",
|
|
help="Move new files to import dir (default: record-only)")
|
|
|
|
args = parser.parse_args()
|
|
cfg = load_config()
|
|
setup_logging(cfg["log_level"])
|
|
|
|
if args.command == "scan":
|
|
scan(cfg, args.path, args.do_import)
|
|
else:
|
|
daemon(cfg)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|