Aberrant/rtorrent.py

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

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)


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 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=',
	)
	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 = xmlrpc.client.ServerProxy(f"http://localhost:8000/RPC{n}")
		sp.append(s)
	WATCH_HANDLE = Watch()
	WATCH_HANDLE.start()

def stop_watch(sig, frame):
	"""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 = [t.tracker for t in all_torrents()]
	tracker_stats = {tr: trackers.count(tr) for tr in set(trackers)}
	tracker_stats = dict(sorted(tracker_stats.items()))
	tracker_stats['total'] = len(trackers)
	return tracker_stats