From 7d2c365c57a8f0f226485aea50b5118d1c4f4232 Mon Sep 17 00:00:00 2001 From: Jeffrey Serio <23226432+hyperreal64@users.noreply.github.com> Date: Sat, 10 Aug 2024 17:34:03 -0500 Subject: [PATCH] feat: make it fancy --- bin/update_tracker.py | 153 ++++++++++++++++++++++++++---------------- 1 file changed, 95 insertions(+), 58 deletions(-) diff --git a/bin/update_tracker.py b/bin/update_tracker.py index 674dbfa..487bdbd 100755 --- a/bin/update_tracker.py +++ b/bin/update_tracker.py @@ -12,11 +12,12 @@ Expectations: - rsync installed on the host system running this script Usage: - update_tracker.py + update_tracker.py [--add-nirn-tracker] update_tracker.py -h Options: - -h, --help show this help message and exit + --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 @@ -26,72 +27,108 @@ 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 - # JSON file containing authentication info for each qBittorrent instance - QBITTORRENT_AUTH_FILE = Path.home().joinpath(".config/qbittorrent_auth.json") + 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) + # 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 - torrent_infohashes = [] - for item in auth_data["instances"]: - qb = Client(item["hostname"]) - qb.login(username=item["username"], password=item["password"]) + # 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 + for torrent in qb.torrents(): + torrent_infohashes.append(torrent.get("hash")) # type: ignore - # Format the infohashes to have a \n at the end - format_infohashes = set([f"{infohash}\n" for infohash in torrent_infohashes]) + # 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 - with tempfile.NamedTemporaryFile() as ntf: - with open(ntf.name, "w") as tf: - tf.writelines(format_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 - subprocess.run( - [ - "scp", - ntf.name, - "root@bttracker.nirn.quest:/etc/opentracker/whitelist.txt", - ] + # 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" ) - # Use SSH to run `systemctl restart opentracker.service` on the remote - # torrent tracker server - subprocess.run( - [ - "ssh", - "root@bttracker.nirn.quest", - "systemctl", - "restart", - "opentracker.service", - ] - ) + console.print(":white_check_mark: reannounce all torrents to their trackers") - # Ensure 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", - ) - print(f"Updated {torrent.get('name')}") # type: ignore - - # Reannounce all torrents in each qBittorrent instance 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) + torrents = Text(str(len(torrent_infohashes))) + torrents.stylize("bold green") + console.print(torrents + " torrents were updated")