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 = speedrate(self.downrate)
		self.uprate = raw[5]
		self.uprate_str = speedrate(self.uprate)

		self.tracker = get_tracker(raw[6])


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 speedrate(rate):
	"""Helper to assign appropriate prefixes to speed rates."""
	unit = "B/s"
	if rate > 1024:
		rate /= 1024
		unit = "KiB/s"
	if rate > 1024:
		rate /= 1024
		unit = "MiB/s"
	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 open_proxy():
	"""Opens connections to all of the rtorrent instances."""
	global sp
	for n in range(NUM_INST):
		s = xmlrpc.client.ServerProxy(f"http://localhost:8000/RPC{n}")
		sp.append(s)

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=',
	)
	return [Torrent(raw) for raw in res]

def init():
	"""Initializes the rTorrent interface."""
	global WATCH_HANDLE
	open_proxy()
	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