admin-scripts/bin/update_tracker.py

#!/usr/bin/env python3

"""update_tracker.py

Description:
This script collects infohashes of all torrents in each qBittorrent instance,
updates opentracker, and reannounces all torrents to their trackers.

Expectations:
- A JSON qBittorrent authentication file at ~/.config/qbittorrent_auth.json
- SSH pubkey access to torrent tracker server root@bttracker.nirn.quest
- rsync installed on the host system running this script

Usage:
    update_tracker.py [--add-nirn-tracker]
    update_tracker.py -h

Options:
    --add-nirn-tracker      ensure bttracker.nirn.quest:6969/announce is added to each torrent's tracker list
    -h, --help              show this help message and exit
"""

import json
import subprocess
import tempfile
from pathlib import Path

from docopt import docopt
from qbittorrent import Client
from rich.console import Console
from rich.text import Text

if __name__ == "__main__":
    args = docopt(__doc__)  # type: ignore

    console = Console()
    with console.status("[bold green]Executing the tasks...") as status:
        # JSON file containing authentication info for each qBittorrent instance
        QBITTORRENT_AUTH_FILE = Path.home().joinpath(".config/qbittorrent_auth.json")

        # Open authentication file and load JSON data
        with open(QBITTORRENT_AUTH_FILE, "r") as qbt_auth:
            auth_data = json.load(qbt_auth)

        # Collect infohashes of all torrents in each qBittorrent instance
        console.log(
            "Collecting infohashes of all torrents in each qBittorrent instance"
        )
        torrent_infohashes = []
        for item in auth_data["instances"]:
            qb = Client(item["hostname"])
            qb.login(username=item["username"], password=item["password"])

            for torrent in qb.torrents():
                torrent_infohashes.append(torrent.get("hash"))  # type: ignore

        # Format the infohashes to have a \n at the end
        console.log("Formatting infohashes to have a newline at the end")
        format_infohashes = set([f"{infohash}\n" for infohash in torrent_infohashes])

        # Create a NamedTemporaryFile and write all infohashes to it, one per line
        console.log("Creating temporary file to write infohashes to")
        with tempfile.NamedTemporaryFile() as ntf:
            with open(ntf.name, "w") as tf:
                tf.writelines(format_infohashes)

            # Use scp to copy the infohashes file to the torrent tracker's config
            # directory on the remote torrent tracker server, overwriting the
            # whitelist.txt file
            console.log(
                "SSH-copying the temporary infohashes file to the torrent tracker's whitelist"
            )
            subprocess.run(
                [
                    "scp",
                    "-q",
                    ntf.name,
                    "root@bttracker.nirn.quest:/etc/opentracker/whitelist.txt",
                ]
            )

            # Use SSH to run `systemctl restart opentracker.service` on the remote
            # torrent tracker server
            console.log("Restarting opentracker.service on the remote server")
            subprocess.run(
                [
                    "ssh",
                    "root@bttracker.nirn.quest",
                    "systemctl",
                    "restart",
                    "opentracker.service",
                ]
            )

        # Ensure bttracker.nirn.quest:6969/announce is added to each torrent's
        # tracker list.
        if args["--add-nirn-tracker"]:
            console.log(
                "Ensuring bttracker.nirn.quest:6969/announce is added to each torrent's tracker list"
            )
            for item in auth_data["instances"]:
                qb = Client(item["hostname"])
                qb.login(username=item["username"], password=item["password"])
                for torrent in qb.torrents():
                    qb.add_trackers(
                        torrent.get("hash"),  # type: ignore
                        "http://bttracker.nirn.quest:6969/announce\nudp://bttracker.nirn.quest:6969/announce",
                    )

        # Reannounce all torrents in each qBittorrent instance to their trackers
        console.log("Reannouncing all torrents to their trackers")
        for item in auth_data["instances"]:
            qb = Client(item["hostname"])
            qb.login(username=item["username"], password=item["password"])
            torrent_infohashes = [torrent.get("hash") for torrent in qb.torrents()]  # type: ignore
            qb.reannounce(torrent_infohashes)

    # Print output and make it look sexy ;)
    console = Console()
    tasks = Text("Tasks completed:\n")
    tasks.stylize("bold magenta")
    console.print(tasks)
    console.print(":white_check_mark: update the tracker's whitelist")

    if args["--add-nirn-tracker"]:
        console.print(
            ":white_check_mark: ensure bttracker.nirn.quest:6969/announce is in each torrent's tracker list"
        )

    console.print(":white_check_mark: reannounce all torrents to their trackers")

    torrents = Text(str(len(torrent_infohashes)))
    torrents.stylize("bold green")
    console.print(torrents + " torrents were updated")