torrent-audit: tag torrents in qbt

services/arr/torrent-audit.nix

@@ -35,6 +35,7 @@
       SONARR_URL = "http://localhost:${builtins.toString service_configs.ports.private.sonarr.port}";
       SONARR_CONFIG = "${service_configs.sonarr.dataDir}/config.xml";
       CATEGORIES = "tvshows,movies,anime";
+      TAG_TORRENTS = "true";
     };
   };
 }

services/arr/torrent-audit.py

@@ -266,6 +266,51 @@ def print_section(torrents, show_status=False):
     print()
 
 
+AUDIT_TAGS = {"audit:unmanaged", "audit:abandoned-safe", "audit:abandoned-review"}
+
+
+def tag_torrents(qbit_client, qbit_torrents, all_known, all_abandoned):
+    log.info("Tagging torrents ...")
+
+    abandoned_by_hash = {t["hash"].upper(): t for t in all_abandoned}
+
+    all_hashes = []
+    for torrents in qbit_torrents.values():
+        all_hashes.extend(torrents.keys())
+
+    for h in all_hashes:
+        current_tags = set()
+        torrent_info = None
+        for torrents in qbit_torrents.values():
+            if h in torrents:
+                torrent_info = torrents[h]
+                break
+        if not torrent_info:
+            continue
+
+        existing_tags = {t.strip() for t in torrent_info.get("tags", "").split(",") if t.strip()}
+        existing_audit_tags = existing_tags & AUDIT_TAGS
+
+        if h in abandoned_by_hash:
+            status = abandoned_by_hash[h]["status"]
+            desired = "audit:abandoned-safe" if status == "SAFE" else "audit:abandoned-review"
+        elif h not in all_known:
+            desired = "audit:unmanaged"
+        else:
+            desired = None
+
+        tags_to_remove = existing_audit_tags - ({desired} if desired else set())
+        tags_to_add = ({desired} if desired else set()) - existing_audit_tags
+
+        low_hash = torrent_info["hash"]
+        for tag in tags_to_remove:
+            qbit_client.torrents_remove_tags(tags=tag, torrent_hashes=low_hash)
+        for tag in tags_to_add:
+            qbit_client.torrents_add_tags(tags=tag, torrent_hashes=low_hash)
+
+    log.info("Tagging complete")
+
+
 def main():
     qbit_url = os.environ["QBITTORRENT_URL"]
     radarr_url = os.environ["RADARR_URL"]
@@ -328,6 +373,10 @@ def main():
     )
     print(f"SAFE TO RECLAIM: {gib(sum(t['size'] for t in safe))} GiB")
 
+    # -- Tagging --
+    if os.environ.get("TAG_TORRENTS", "").lower() in ("1", "true", "yes"):
+        tag_torrents(qbit, qbit_torrents, all_known, all_abandoned)
+
 
 if __name__ == "__main__":
     main()