Aberrant/rtorrent.py

#!/usr/bin/env python3
"""
This module handles the interface with rTorrent via XMLRPC.
"""
import re
import time
import threading
from collections import defaultdict

import rtorrent_xmlrpc

NUM_INST = 10
WATCH_HANDLE = None
sp = []
torrents = [[]] * NUM_INST

class Torrent:
	def __init__(self, raw):
		self.hash = raw[0]
		self.name = raw[1]

		self.active = raw[2]
		self.complete = raw[3]
		if not self.active:
			self.state = "inactive"
		elif self.complete:
			self.state = "seeding"
		else:
			self.state = "leeching"

		self.downrate = raw[4]
		self.downrate_str = size_units(self.downrate) + '/s'
		self.uprate = raw[5]
		self.uprate_str = size_units(self.uprate) + '/s'

		self.tracker = get_tracker(raw[6])

		self.down_total = raw[7]
		self.total_size = raw[8]
		self.total_size_str = size_units(self.total_size)
		self.down_percent = round((self.down_total / self.total_size) * 100, 2)
		if self.state == "leeching":
			if self.downrate:
				self.eta = (self.total_size - self.down_total) / self.downrate
				self.eta_str = time_units(self.eta)
			else:
				self.eta = float("inf")
				self.eta_str = "∞"
		else:
			self.eta = 0
			self.eta_str = ""

		self.message = raw[9]
		self.hashing = raw[10]


class Watch(threading.Thread):
	"""A thread class that continously queries the rTorrent instances."""
	def __init__(self):
		super(Watch, self).__init__()
		self._stop_event = threading.Event()

	def stop(self):
		self._stop_event.set()

	def stopped(self):
		return self._stop_event.is_set()

	def run(self):
		global torrents
		while not self.stopped():
			for n in range(NUM_INST):
				if self.stopped():
					break
				torrents[n] = get_all(n)
				self._stop_event.wait(2)


def size_units(rate):
	"""Helper to assign appropriate prefixes to numbers."""
	unit = "B"
	if rate > 1024:
		rate /= 1024
		unit = "KiB"
	if rate > 1024:
		rate /= 1024
		unit = "MiB"
	if rate > 1024:
		rate /= 1024
		unit = "GiB"
	rate = round(rate, 1)
	return str(rate) + unit

def time_units(seconds):
	"""Helper to convert seconds into more useful units."""
	seconds = int(seconds)
	if seconds > (24*60*60):
		days = seconds // (24*60*60)
		hours = (seconds % (24*60*60)) // (60*60)
		eta = f"{days}d{hours}h"
	elif seconds > (60*60):
		eta = f"{hours}h{minutes}m"
	elif seconds > 60:
		minutes = seconds // 60
		seconds = seconds % 60
		eta = f"{minutes}m{seconds}s"
	else:
		eta = f"{seconds}s"
	return eta

def get_tracker(path):
	"""
	At present I don't have an efficient way to get the tracker url
	with the d.multicall2() function, so we parse it from the
	directory path.
	"""
	return path.split('/')[4]

def all_torrents():
	"""Helper that returns a list of all torrents."""
	res = []
	for item in torrents:
		res += item
	return res

def get_all(n):
	"""Gets all torrent information from a instance and returns it."""
	res = sp[n].d.multicall2('', 'main',
		'd.hash=',
		'd.name=',
		'd.is_active=',
		'd.complete=',
		'd.down.rate=',
		'd.up.rate=',
		'd.directory=',
		'd.completed_bytes=',
		'd.size_bytes=',
		'd.message=',
		'd.hashing=',
	)
	return [Torrent(raw) for raw in res]

def init():
	"""Initializes the rTorrent interface."""
	global WATCH_HANDLE
	global sp
	for n in range(NUM_INST):
		s = rtorrent_xmlrpc.SCGIServerProxy(f"scgi://localhost:500{n}")
		sp.append(s)
	WATCH_HANDLE = Watch()
	WATCH_HANDLE.start()

def stop_watch(*args, **kwargs):
	"""Stops the watch thread."""
	global WATCH_HANDLE
	WATCH_HANDLE.stop()

def get_active():
	"""Returns all actively seeding or leeching torrents."""
	active = [t for t in all_torrents() if t.downrate or t.uprate]
	return active

def get_stats():
	"""Returns various statistical information about the torrents."""
	trackers = {}
	for torrent in all_torrents():
		stats = [0]*3
		stats[0] = 1 if torrent.hashing else 0
		stats[1] = 1 if torrent.message else 0
		stats[2] = 1

		if trackers.get(torrent.tracker):
			for i in range(len(stats)):
				trackers[torrent.tracker][i] += stats[i]
		else:
			trackers[torrent.tracker] = stats
	trackers = dict(sorted(trackers.items()))
	total = [0]*3
	for i in range(len(total)):
		total[i] = sum([v[i] for _, v in trackers.items()])
	trackers['total'] = total
	return trackers