OppaiBot/oppaibot.py

#!/usr/bin/env python3
"""
A lazy bot for idling in OppaiTime's announce channel and downloading
every torrent with matching tags.
"""
import os
import re
import functools
import configparser
import http.cookiejar

import requests
from twisted.internet import protocol, reactor
from twisted.words.protocols import irc

HEADERS = {"User-Agent": "spaghetti is a faggot"}


class OppaiBot(irc.IRCClient):
	def __init__(self, config):
		self.config = config
		self.tracker = self.config.name
		self.nickname = self.config["nickname"]
		self.username = self.config["ident"]

		self.cj = http.cookiejar.MozillaCookieJar(self.config["cookies_txt"])
		self.cj.load()

		self.watch_dir = self.config["watch_dir"]
		
		self.tags = re.findall(r"\[(.+?)\]", self.config["tags"])
		self.tags = [tag.split(",") for tag in self.tags]
		self.blacklist = self.config.get("blacklist", "").split(",")


	def save_torrent(self, url, directory):
		"""
		Downloads and saves the torrent file to the given directory.
		"""
		res = requests.get(url, cookies=self.cj, headers=HEADERS, verify=True)
		res.raise_for_status()
		fname =re.search(r'filename="(.+)"',res.headers['content-disposition'])
		fname = fname.group(1)
		fname = fname.encode("latin-1").decode("utf-8")
		print(self.tracker + ": Saving torrent:", fname)
		with open(os.path.join(directory, fname), "wb") as file:
			for chunk in res.iter_content(100000):
				file.write(chunk)


	def privmsg(self, user, channel, message):
		"""
		Called when the bot receives a PRIVMSG, which can come from channels
		or users alike.
		"""
		# More advanced logic is left as an exercise to the reader.
		black_in_msg = [tag in message for tag in self.blacklist if tag]
		if any(black_in_msg):
			return

		tags_in_msg = []
		for tag_group in self.tags:
			tags_in_group = []
			for tag in tag_group:
				if tag.startswith("!"):
					tags_in_group.append(tag[1:] not in message)
				else:
					tags_in_group.append(tag in message)
			tags_in_msg.append(all(tags_in_group))

		if any(tags_in_msg):
			url = re.findall(r"(http.+?) ", message)[1]
			self.save_torrent(url, self.watch_dir)


	def joined(self, channel):
		"""Called when the bot joins a new channel."""
		print(self.tracker + ": Joined", channel)


	def signedOn(self):
		"""Called when the bot successfully connects to the server."""
		print(self.tracker + ": Signed on as", self.nickname)
		self.mode(self.nickname, True, "B") # set +B on self
		self.msg(self.config["tracker_bot"], self.config["auth_msg"])


	def nickChanged(self, nick):
		"""Called when my nick has been changed."""
		print(self.tracker + ": Nick changed to", nick)


class OppaiBotFactory(protocol.ReconnectingClientFactory):
	# black magic going on here
	protocol = property(lambda s: functools.partial(OppaiBot, s.config))

	def __init__(self, config):
		self.config = config


if __name__ == "__main__":
	import argparse

	parser = argparse.ArgumentParser(
		description="Downloads new torrents from a torrent tracker's IRC " \
			+ "announce channel.")
	parser.add_argument(
		"-c",
		"--config",
		default="config.cfg",
		help="Specify an alternate config file to use.")
	args = parser.parse_args()

	config = configparser.ConfigParser()
	config.read(args.config)
	for tracker in config.sections():
		tracker_config = config[tracker]

		server = tracker_config["server"]
		port = tracker_config.getint("port")
		print("Connecting to:", tracker_config.name)
		reactor.connectTCP(server, port, OppaiBotFactory(tracker_config))
	reactor.run()