From 9c6ae886cebbc39dec02752cd2f894d471758c02 Mon Sep 17 00:00:00 2001 From: Nikolai R Kristiansen Date: Tue, 22 Mar 2016 13:51:00 +0100 Subject: [PATCH] Spring cleaning --- .gitignore | 3 +- README | 4 +- requirements.txt | 10 +- settings-sample.py | 7 -- snurr.py | 241 +++++++++++++++------------------------------ tests.py | 25 +++++ utils.py | 74 ++++++++++++++ 7 files changed, 189 insertions(+), 175 deletions(-) create mode 100644 tests.py create mode 100644 utils.py diff --git a/.gitignore b/.gitignore index fbe1365..904200b 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ settings.py *.swp *.pyc -venv +/venv +/.idea \ No newline at end of file diff --git a/README b/README index f960e9b..5264e1a 100644 --- a/README +++ b/README @@ -1,12 +1,12 @@ # Requirements: With virtualenv: - apt-get install libxml2-dev libxslt1-dev python-dev libmysqlclient-dev + apt install libxml2-dev libxslt1-dev python-dev libmysqlclient-dev virtualenv venv . venv/bin/activate pip install -r requirements.txt Without venv: - python-twisted python-mysqldb python-lxml python-openssl + python-twisted python-mysqldb python-lxml python-openssl python-requests $ python snurr.py -h Usage: python snurr.py [options] CHANNEL diff --git a/requirements.txt b/requirements.txt index bcad18d..3825ddc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,4 @@ -Twisted==15.0.0 -lxml==3.4.2 -mysqlclient==1.3.6 -pyOpenSSL==0.14 -service-identity==14.0.0 - +Twisted[tls]==16.0.0 +lxml==3.6.0 +mysqlclient==1.3.7 +requests==2.9.1 diff --git a/settings-sample.py b/settings-sample.py index 392d2d1..ca896f8 100644 --- a/settings-sample.py +++ b/settings-sample.py @@ -1,12 +1,5 @@ -# snurrbot configuration -DISABLE_LOG = False # DATABASE DB_API_ADAPTER = "MySQLdb" # DB-API 2.0 adapter -DB_HOST = "localhost" # hostname -DB_USER = "snurr" # username -DB_PASSWORD = "" -DB_NAME = "snurr" - DISABLE_TETRIS = False TETRIS_DB_HOST = "localhost" # hostname TETRIS_DB_USER = "snurr_tetris" # username diff --git a/snurr.py b/snurr.py index 309c591..c0222ee 100755 --- a/snurr.py +++ b/snurr.py @@ -1,43 +1,38 @@ #!/usr/bin/env python # encoding: utf-8 import optparse -import subprocess -import time -import MySQLdb -import urllib2, lxml.html, re -from datetime import datetime - - +import requests +import re +import logging from twisted.words.protocols import irc from twisted.internet import protocol, reactor, ssl -from twisted.enterprise import adbapi -from twisted.python import log import settings -if not settings.DISABLE_LOG: - import MySQLdb +from utils import ReconnectingConnectionPool, ping_host, parse_title + +logger = logging.getLogger('snurr') + class SnurrBot(irc.IRCClient): def __init__(self): self.actions = IRCActions(self) - def _get_nickname(self): + @property + def nickname(self): return self.factory.nickname - nickname = property(_get_nickname) - def __unicode__(self): - return "SnurrBot:%s" % (nickname,) + return "SnurrBot:%s" % self.nickname def signedOn(self): self.join(self.factory.channel) - _log("Signed on as %s." % (self.nickname,)) + logger.info("Signed on as %s." % (self.nickname,)) # make the client instance known to the factory self.factory.bot = self def joined(self, channel): - _log("Joined %s." % (channel,)) + logger.info("Joined %s." % channel) # Called when I have a message from a user to me or a channel. def privmsg(self, user, channel, msg): @@ -46,22 +41,29 @@ class SnurrBot(irc.IRCClient): self.actions.new(msg[1:], user, channel) else: self.actions.newfull(msg, user, channel) - _log("PRIVMSG: %s: %s" % (user,msg,)) + logger.info("PRIVMSG: %s: %s" % (user,msg)) def msgReply(self, user, to, msg): if len(msg) > 0: if to == self.nickname: - _log("Message sent to %s" % (user,)) + logger.info("Message sent to %s" % user) self.msg(user, msg, length=512) else: self.msgToChannel(msg) def msgToChannel(self, msg): # Sends a message to the predefined channel - _log("Message sent to %s" % (self.factory.channel,)) + logger.info("Message sent to %s" % self.factory.channel) + if len(msg) > 0: self.say(self.factory.channel, msg, length=512) + def rawDataReceived(self, data): + pass + + def dccSend(self, user, file): + pass + class SnurrBotFactory(protocol.ClientFactory): protocol = SnurrBot @@ -74,11 +76,12 @@ class SnurrBotFactory(protocol.ClientFactory): return "SnurrBotFactory" def clientConnectionLost(self, connector, reason): - _log("Lost connection (%s), reconnecting." % (reason,)) + logger.info("Lost connection (%s), reconnecting." % reason) connector.connect() def clientConnectionFailed(self, connector, reason): - _log("Could not connect: (%s), retrying." % (reason,)) + logger.info("Could not connect: (%s), retrying." % reason) + class UDPListener(protocol.DatagramProtocol): @@ -86,102 +89,70 @@ class UDPListener(protocol.DatagramProtocol): self.botfactory = botfactory def startProtocol(self): - _log("Listening for messages.") + logger.info("Listening for messages.") def stopProtocol(self): - _log("Listener stopped") + logger.info("Listener stopped") - def datagramReceived(self, data, (host, port)): + def datagramReceived(self, data, addr): + host, port = addr # UDP messages (from MediaWiki f.ex) arrive here and are relayed # to the ircbot created by the bot factory. - _log("Received %r from %s:%d" % (data, host, port)) - _log("Relaying msg to IRCClient in %s" % (self.botfactory,)) + logger.info("Received %r from %s:%d" % (data, host, port)) + logger.info("Relaying msg to IRCClient in %s" % (self.botfactory,)) if self.botfactory.bot: self.botfactory.bot.msgToChannel(data) -class IRCActions(): + +class IRCActions: def __init__(self, bot): self.bot = bot - if not settings.DISABLE_LOG: - self.dbpool = self._get_dbpool() if not settings.DISABLE_TETRIS: self.tetris_dbpool = self._get_tetris_dbpool() - def _get_dbpool(self): - # Setup an async db connection - return ReconnectingConnectionPool(settings.DB_API_ADAPTER, - host=settings.DB_HOST, user=settings.DB_USER, - passwd=settings.DB_PASSWORD, db=settings.DB_NAME, - charset="utf8", use_unicode=True, - cp_reconnect=True) - def _get_tetris_dbpool(self): # Setup an async db connection - return ReconnectingConnectionPool(settings.DB_API_ADAPTER, + return ReconnectingConnectionPool( + settings.DB_API_ADAPTER, host=settings.TETRIS_DB_HOST, user=settings.TETRIS_DB_USER, passwd=settings.TETRIS_DB_PASSWORD, db=settings.TETRIS_DB_NAME, charset="utf8", use_unicode=True, cp_reconnect=True) - def ping(self, host): - # TODO rewrite async - try: - command = "ping -W 1 -c 1 " + host - retcode = subprocess.call(command.split(),stdout=subprocess.PIPE) - if retcode == 0: - return host + " pinger fint den :P" - elif retcode == 2: - return host + " pinger ikke :(" - else: - return "ping returned: " + str(retcode) - except OSError, e: - _log("Execution failed:" + str(e)) - return "feil med ping:" + str(e) - def help(self): text = "" text += "Command: !help\n" text += " This help message\n" - if not settings.DISABLE_LOG: - text += "Command: !log DESCRIPTION\n" - text += " Add new entry in log\n" - text += "Command: !lastlog\n" - text += " Last 3 log entries\n" + if not settings.DISABLE_TETRIS: text += "Command: !tetrishigh\n" text += " Display tetris highscore\n" text += "Command: !ping HOST\n" text += " Ping target host" + return text def new(self, msg, user, channel): - nick = user.split('!', 1)[0] # Strip out hostmask + nick = user.split('!', 1)[0] # Strip out hostmask + help_reply = "Need !help " + nick + "?" # Process the commands parts = msg.split() if not parts: - self.bot.msgReply(nick, channel, "Need !help " + nick + "?") + self.bot.msgReply(nick, channel, help_reply) return if parts[0] == "ping" and len(parts) == 2: - self.bot.msgReply(nick, channel, self.ping(parts[1])) + self.bot.msgReply(nick, channel, ping_host(parts[1])) elif parts[0] == "help" and len(parts) == 1: self.bot.msgReply(nick, channel, self.help()) - elif parts[0] == "log" and len(parts) >= 2 and not settings.DISABLE_LOG: - # set_log_entry should create a deferred and - # the callback should fire when the db returns. - self.set_log_entry(nick, parts[1:]).addCallback(self.msg_log_entry, channel, nick) - elif parts[0] == "lastlog" and len(parts) == 1 and not settings.DISABLE_LOG: - # ...same as above - self.get_lastlog().addCallback(self.msg_lastlog, channel, nick) elif parts[0] == "tetrishigh" and len(parts) == 1 and not settings.DISABLE_TETRIS: self.get_tetris_highscore().addCallback(self.msg_tetris_highscore, channel, nick) - else: - self.bot.msgReply(nick, channel, "Need !help " + nick + "?") + self.bot.msgReply(nick, channel, help_reply) def newfull(self, msg, user, channel): - nick = user.split('!', 1)[0] # Strip out hostmask + nick = user.split('!', 1)[0] # Strip out hostmask # Check if msg contains HTTP(S) URL urlmatch = re.search(r"(https?):\/\/([\da-z\.-]+)\.([a-z\.]{2,6})([\/\w\.-]*\/?)(\?[\/\w=&\.-]+)?", msg, re.I) @@ -192,122 +163,74 @@ class IRCActions(): reply = "URL: " url = urlmatch.group() try: - req = urllib2.Request(url, headers={'User-Agent': 'snurrbot v0.1'}) - i = urllib2.urlopen(req) + req = requests.get(url, headers={'User-Agent': 'snurrbot v0.1'}, stream=True) except Exception as e: - _log(e) + logger.info(e) self.bot.msgReply(nick, channel, reply + "Timeout") return - statuscode = i.getcode() - contenttype = i.info().gettype() - reply += "%d %s " % (statuscode, contenttype) - if statuscode == 200 and contenttype == "text/html": - # Read max 10000 bytes - partialhtml = i.read(10000) - try: - # This might raise an exception if the HTML is _really_ broken - tree = lxml.html.fromstring(partialhtml) - title = tree.find(".//title") - if title is not None and title.text is not None: - reply += "[%s]" % re.sub(r'\s+', ' ', title.text.strip()) - # DEBUG - if title is None: - reply += "" - elif title.text is None: - reply += "<title.text is none>" - # END - except: - reply += "<exception>" - pass + + content_type = req.headers['content-type'] + if ';' in content_type: + content_type = content_type.split(';')[0] + reply += "%d %s " % (req.status_code, content_type) + if req.status_code == 200 and content_type == "text/html": + # Read max 10240 bytes + partial_html = next(req.iter_content(chunk_size=10240)) + reply += parse_title(partial_html) + self.bot.msgReply(nick, channel, reply) - def set_log_entry(self, nick, entry): - entry = " ".join(entry) - sql = "INSERT INTO main_entry (user, text, created) VALUES (%s, %s, NOW())" - params = (nick, entry) - return self.dbpool.runOperation(sql, params) - - def msg_log_entry(self, result, channel, nick): - self.bot.msgReply(nick, channel, "Yes sir! Driftslogg oppdatert.") - - def get_lastlog(self): - sql = "SELECT * FROM main_entry ORDER BY created DESC LIMIT 3" - return self.dbpool.runQuery(sql) - - def msg_lastlog(self, log_entries, channel, nick): - for i,entry in enumerate(log_entries, start=1): - string_entry = str(i) + ": " + entry[2] + " (" + entry[1] + ", " + str(entry[3]) + ")" - self.bot.msgReply(nick, channel, string_entry.encode("utf-8")) - def get_tetris_highscore(self): sql = "SELECT MAX(score) AS highscore, name FROM highscore GROUP BY name ORDER BY highscore DESC LIMIT 3" return self.tetris_dbpool.runQuery(sql) def msg_tetris_highscore(self, highscores, channel, nick): for highscore in highscores: - high,name = highscore + high, name = highscore string_highscore = "Highscore: " + str(high) + " by " + name + "." self.bot.msgReply(nick, channel, string_highscore.encode("utf-8")) -class ReconnectingConnectionPool(adbapi.ConnectionPool): - """Reconnecting adbapi connection pool for MySQL. - This class improves on the solution posted at - http://www.gelens.org/2008/09/12/reinitializing-twisted-connectionpool/ - by checking exceptions by error code and only disconnecting the current - connection instead of all of them. - - Also see: - http://twistedmatrix.com/pipermail/twisted-python/2009-July/020007.html - - """ - def _runInteraction(self, interaction, *args, **kw): - try: - return adbapi.ConnectionPool._runInteraction(self, interaction, *args, **kw) - except MySQLdb.OperationalError, e: - if e[0] not in (2006, 2013): - raise - log.msg("RCP: got error %s, retrying operation" %(e)) - conn = self.connections.get(self.threadID()) - self.disconnect(conn) - # try the interaction again - return adbapi.ConnectionPool._runInteraction(self, interaction, *args, **kw) - -def _log(message): - now = datetime.now().strftime("%b %d %H:%M:%S") - print "{0} {1}".format(now, message) - -def _setup_and_parse_options(): +def _get_parser(): + usage = 'python snurr.py [-h] [options] CHANNEL' parser = optparse.OptionParser(description='Pipes UDP-messages to an IRC-channel.', - usage=_usage()) + usage=usage) parser.add_option('-c', '--connect', metavar='SERVER', - help='IRC server (default: irc.ifi.uio.no)', default='irc.oftc.net') + help='IRC server (default: irc.oftc.net)', default='irc.oftc.net') parser.add_option('-p', '--port', metavar='PORT', type=int, help='IRC server port (default: 6697)', default=6697) - parser.add_option('-s','--ssl', action='store_true', + parser.add_option('-s', '--ssl', action='store_true', help='connect with SSL (default: False)', default=True) parser.add_option('-l', '--listen_port', metavar='LISTEN_PORT', type=int, help='UDP listen port (default: 55666)', default=55666) - return parser.parse_args() + return parser -def _usage(): - return 'Usage: python snurr.py [-h] [options] CHANNEL' if __name__ == "__main__": - options, args = _setup_and_parse_options() + LOG_FORMAT = '[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s' + logger.setLevel(logging.DEBUG) + consoleLogger = logging.StreamHandler() + consoleLogger.setFormatter(logging.Formatter(LOG_FORMAT)) + logger.addHandler(consoleLogger) + + argument_parser = _get_parser() + options, args = argument_parser.parse_args() if len(args) != 1: - print _usage() - exit() - snurr = SnurrBotFactory('#' + args[0]) - listener = UDPListener(snurr) - - # Start the listener. - reactor.listenUDP(options.listen_port, listener) + print(argument_parser.usage) + exit(1) + channel = args[0] if args[0][0] == '#' else '#' + args[0] # Start IRC-bot on specified server and port. + my_snurrbot = SnurrBotFactory(channel) if options.ssl: - reactor.connectSSL(options.connect, options.port, snurr, ssl.ClientContextFactory()) + reactor.connectSSL(options.connect, options.port, my_snurrbot, ssl.ClientContextFactory()) else: - reactor.connectTCP(options.connect, options.port, snurr) + reactor.connectTCP(options.connect, options.port, my_snurrbot) + + # Start the MediaWiki changelog listener. + mediawiki_listener = UDPListener(my_snurrbot) + reactor.listenUDP(options.listen_port, mediawiki_listener) + + # Fire up the reactor reactor.run() diff --git a/tests.py b/tests.py new file mode 100644 index 0000000..0ec0586 --- /dev/null +++ b/tests.py @@ -0,0 +1,25 @@ +import unittest + +import requests +from utils import parse_title + + +class ParseHTMLTest(unittest.TestCase): + def test_parse_title(self): + title = 'Adventure Time GIF - Find Share on GIPHY' + partial_html = """<!DOCTYPE html><html itemscope itemtype="http://schema.org/WebPage" > + <head > + <title>{} + + """.format(title) + self.assertEqual(parse_title(partial_html=partial_html), '[{}]'.format(title)) + + # @unittest.skip() + def test_parse_title_from_stream(self): + url = 'http://lol.com' + req = requests.get(url, headers={'User-Agent': 'snurrbot v0.1'}, stream=True) + partial_html = req.iter_content(chunk_size=10240) + self.assertEqual(parse_title(partial_html=partial_html), None) + +if __name__ == '__main__': + unittest.main() diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..99cd18c --- /dev/null +++ b/utils.py @@ -0,0 +1,74 @@ +import logging +import re +import subprocess + +import lxml.html +from twisted.enterprise import adbapi + + +logger = logging.getLogger(__name__) + + +class ReconnectingConnectionPool(adbapi.ConnectionPool): + """Reconnecting adbapi connection pool for MySQL. + + This class improves on the solution posted at + http://www.gelens.org/2008/09/12/reinitializing-twisted-connectionpool/ + by checking exceptions by error code and only disconnecting the current + connection instead of all of them. + + Also see: + http://twistedmatrix.com/pipermail/twisted-python/2009-July/020007.html + + """ + def _runInteraction(self, interaction, *args, **kw): + import MySQLdb + + try: + return adbapi.ConnectionPool._runInteraction(self, interaction, *args, **kw) + except MySQLdb.OperationalError as e: + if e[0] not in (2006, 2013): + raise + + logger.info("RCP: got error %s, retrying operation" % e) + conn = self.connections.get(self.threadID()) + self.disconnect(conn) + # try the interaction again + return adbapi.ConnectionPool._runInteraction(self, interaction, *args, **kw) + + +def ping_host(host): + # FIXME rewrite async + try: + command = "ping -W 1 -c 1 " + host + retcode = subprocess.call(command.split(), stdout=subprocess.PIPE) + if retcode == 0: + return host + " pinger fint den :P" + elif retcode == 2: + return host + " pinger ikke :(" + else: + return "ping returned: " + str(retcode) + except OSError as e: + logger.info("Execution failed:" + str(e)) + return "feil med ping:" + str(e) + + +def parse_title(partial_html): + reply = '' + try: + # This might raise an exception if the HTML is _really_ broken + tree = lxml.html.fromstring(partial_html) + title = tree.find(".//title") + if title is not None and title.text is not None: + reply += "[%s]" % re.sub(r'\s+', ' ', title.text.strip()) + # DEBUG + if title is None: + reply += "" + elif title.text is None: + reply += "<title.text is none>" + # END + except Exception as e: + reply += "<exception>" + logger.error(e.message) + + return reply