Aberrant/rtorrent.py

122 lines
2.5 KiB
Python

#!/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