admin-scripts/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
- rsync installed on the host system running this script

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

Options:
    --add-tracker DOMAIN           ensure the provided tracker domain is added to each torrent's tracker list
    -h, --help                     show this help message and exit

Examples:
    update_tracker.py --add-tracker hyperreal.coffee
    update_tracker.py
"""

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

    tracker_domain = args["--add-tracker"]

    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,
                    f"root@{tracker_domain}:/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",
                    f"root@{tracker_domain}",
                    "systemctl",
                    "restart",
                    "opentracker.service",
                ]
            )

        # Ensure {tracker_domain}:6969/announce is added to each torrent's
        # tracker list.
        if tracker_domain:
            console.log(
                f"Ensuring {tracker_domain}: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
                        f"http://{tracker_domain}:6969/announce\nudp://{tracker_domain}: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)

        console.log("Done!")

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

    if tracker_domain:
        console.print(
            f":white_check_mark: ensure {tracker_domain}: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")