Aberrant/rtorrent_proxy.py

#!/usr/bin/env python3
"""
This file handles a rather silly use case:
You have two servers, A and B, each with various rtorrent instances running.
You want one instance of Aberrant, running on server A. In your infinite
wisdom, you bound the SCGI ports of the rtorrent instances on server B to
localhost, leaving Aberrant unable to query them. Rtorrent won't allow you
to rebind the SCGI port while it's running, and you don't want to restart
the rtorrent instance(s) because restarting stuff sucks. Therefore, you
need another way for Aberrant to communicate to them.

This script acts as a proxy to the rtorrent instances on server B. It will
listen on a port (don't forget to open the port in your iptables), and
relay all data it receives to the torrent SCGI port and vice versa.

To use:
Configure `port_mapping` in the form `listen_port`: `rtorrent_port`.
Run the script on server B.
"""
import socket
import threading

port_mapping = {
	4000: 5000,
	4001: 5001,
	4002: 5002,
}

class SocketProxy(threading.Thread):
	"""Represents a proxy to an rTorrent SCGI port."""
	def __init__(self, listen_port, rtorrent_port):
		super().__init__()
		self.listen_port = listen_port
		self.rtorrent_port = rtorrent_port
		self._stop_event = threading.Event()

	def stop(self):
		"""Sets the internal 'stop running' flag."""
		self._stop_event.set()

	def stopped(self):
		"""Returns true if the internal stop flag has been set."""
		return self._stop_event.is_set()

	def run(self):
		with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
			s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
			s.settimeout(1)
			s.bind(('0.0.0.0', self.listen_port))
			s.listen(1)
			while not self.stopped():
				try:
					conn, addr = s.accept() # client connected
				except socket.timeout:
					continue
				with conn:
					# receive data from client
					data = []
					data.append(conn.recv(1024)) # the header really shouldn't be longer than this...
					header = data[0].partition(b',')[0]
					body_len = int(header.split(b'\x00')[1])
					recv_len = len(data[0])
					while recv_len < body_len + len(header) + 1:
						data.append(conn.recv(1024))
						recv_len += len(data[-1])
					data = b''.join(data)

					with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as rt_conn:
						rt_conn.connect(('127.0.0.1', self.rtorrent_port))
						rt_conn.sendall(data) # send to rtorrent

						# receive data from rtorrent
						data = []
						data.append(rt_conn.recv(1024))
						header = data[0].partition(b'<?xml')[0]
						body_len = header.split(b'\r\n')[2].partition(b' ')[2]
						body_len = int(body_len)
						recv_len = len(data[0])
						while recv_len < body_len + len(header):
							data.append(rt_conn.recv(1024))
							recv_len += len(data[-1])
						data = b''.join(data)
					conn.sendall(data) # send to client

def run_proxies():
	"""Run all proxies."""
	proxies = []
	for listen_port, rtorrent_port in port_mapping.items():
		p = SocketProxy(listen_port, rtorrent_port)
		p.start()
		proxies.append(p)

	try:
		proxies[0].join()
	except KeyboardInterrupt:
		for p in proxies:
			p.stop()

if __name__ == '__main__':
	run_proxies()