diff --git a/irc.py b/irc.py index 1dadc52..a95926f 100755 --- a/irc.py +++ b/irc.py @@ -9,11 +9,16 @@ http://inamidst.com/phenny/ import asynchat import asyncore +import functools +import proto import re import socket import ssl import sys import time +import traceback +import threading +from tools import decorate class Origin(object): @@ -49,9 +54,12 @@ class Bot(asynchat.async_chat): self.channels = channels or [] self.stack = [] - import threading self.sending = threading.RLock() + proto_func = lambda attr: functools.partial(proto.commands[attr], self) + proto_map = {attr: proto_func(attr) for attr in proto.commands} + self.proto = decorate(object(), proto_map) + def initiate_send(self): self.sending.acquire() asynchat.async_chat.initiate_send(self) @@ -61,24 +69,22 @@ class Bot(asynchat.async_chat): # asynchat.async_chat.push(self, *args, **kargs) def __write(self, args, text=None): - # print 'PUSH: %r %r %r' % (self, args, text) - try: - if text is not None: - # 510 because CR and LF count too, as nyuszika7h points out - self.push((b' '.join(args) + b' :' + text)[:510] + b'\r\n') - else: - self.push(b' '.join(args)[:512] + b'\r\n') - except IndexError: - pass + line = b' '.join(args) + + if text is not None: + line += b' :' + text + + # 510 because CR and LF count too + self.push(line[:510] + b'\r\n') def write(self, args, text=None): """This is a safe version of __write""" def safe(input): if type(input) == str: - input = input.replace('\n', '') - input = input.replace('\r', '') + input = re.sub(' ?(\r|\n)+', ' ', input) return input.encode('utf-8') else: + input = re.sub(b' ?(\r|\n)+', b' ', input) return input try: args = [safe(arg) for arg in args] @@ -127,10 +133,12 @@ class Bot(asynchat.async_chat): def handle_connect(self): if self.verbose: print('connected!', file=sys.stderr) + if self.password: - self.write(('PASS', self.password)) - self.write(('NICK', self.nick)) - self.write(('USER', self.user, '+iw', self.nick), self.name) + self.proto.pass_(self.password) + + self.proto.nick(self.nick) + self.proto.user(self.user, '+iw', self.name) def handle_close(self): self.close() @@ -165,7 +173,7 @@ class Bot(asynchat.async_chat): self.dispatch(origin, tuple([text] + args)) if args[0] == 'PING': - self.write(('PONG', text)) + self.proto.pong(text) def dispatch(self, origin, args): pass @@ -203,12 +211,7 @@ class Bot(asynchat.async_chat): self.sending.release() return - def safe(input): - if type(input) == str: - input = input.encode('utf-8') - input = input.replace(b'\n', b'') - return input.replace(b'\r', b'') - self.__write((b'PRIVMSG', safe(recipient)), safe(text)) + self.proto.privmsg(recipient, text) self.stack.append((time.time(), text)) self.stack = self.stack[-10:] @@ -218,12 +221,8 @@ class Bot(asynchat.async_chat): text = "\x01ACTION {0}\x01".format(text) return self.msg(recipient, text) - def notice(self, dest, text): - self.write(('NOTICE', dest), text) - def error(self, origin): try: - import traceback trace = traceback.format_exc() print(trace) lines = list(reversed(trace.splitlines())) diff --git a/modules/admin.py b/modules/admin.py index 5bd1035..aeac5c3 100644 --- a/modules/admin.py +++ b/modules/admin.py @@ -13,9 +13,7 @@ def join(phenny, input): if input.sender.startswith('#'): return if input.admin: channel, key = input.group(1), input.group(2) - if not key: - phenny.write(['JOIN'], channel) - else: phenny.write(['JOIN', channel, key]) + phenny.proto.join(channel, key) join.rule = r'\.join (#\S+)(?: *(\S+))?' join.priority = 'low' join.example = '.join #example or .join #example key' @@ -24,7 +22,7 @@ def autojoin(phenny, input): """Join the specified channel when invited by an admin.""" if input.admin: channel = input.group(1) - phenny.write(['JOIN'], channel) + phenny.proto.join(channel) autojoin.event = 'INVITE' autojoin.rule = r'(.*)' @@ -33,7 +31,7 @@ def part(phenny, input): # Can only be done in privmsg by an admin if input.sender.startswith('#'): return if input.admin: - phenny.write(['PART'], input.group(2)) + phenny.proto.part(input.group(2)) part.rule = (['part'], r'(#\S+)') part.priority = 'low' part.example = '.part #example' @@ -43,7 +41,7 @@ def quit(phenny, input): # Can only be done in privmsg by the owner if input.sender.startswith('#'): return if input.owner: - phenny.write(['QUIT']) + phenny.proto.quit() __import__('os')._exit(0) quit.commands = ['quit'] quit.priority = 'low' diff --git a/modules/startup.py b/modules/startup.py index 0af303b..a77a967 100644 --- a/modules/startup.py +++ b/modules/startup.py @@ -27,13 +27,11 @@ def setup(phenny): timer = threading.Timer(refresh_delay, close, ()) phenny.data['startup.setup.timer'] = timer phenny.data['startup.setup.timer'].start() - # print "PING!" - phenny.write(('PING', phenny.config.host)) + phenny.proto.ping(phenny.config.host) phenny.data['startup.setup.pingloop'] = pingloop def pong(phenny, input): try: - # print "PONG!" phenny.data['startup.setup.timer'].cancel() time.sleep(refresh_delay + 60.0) pingloop() @@ -50,16 +48,16 @@ def startup(phenny, input): if phenny.data.get('startup.setup.pingloop'): phenny.data['startup.setup.pingloop']() - if hasattr(phenny.config, 'serverpass'): - phenny.write(('PASS', phenny.config.serverpass)) + if hasattr(phenny.config, 'serverpass'): + phenny.proto.pass_(phenny.config.serverpass) if hasattr(phenny.config, 'password'): phenny.msg('NickServ', 'IDENTIFY %s' % phenny.config.password) time.sleep(5) # Cf. http://swhack.com/logs/2005-12-05#T19-32-36 - for channel in phenny.channels: - phenny.write(('JOIN', channel)) + for channel in phenny.channels: + phenny.proto.join(channel) time.sleep(0.5) startup.rule = r'(.*)' startup.event = '251' diff --git a/proto.py b/proto.py new file mode 100644 index 0000000..0959ccd --- /dev/null +++ b/proto.py @@ -0,0 +1,55 @@ +#!/usr/bin/env python3 +""" +proto.py - IRC protocol messages +""" + +import sys +import traceback + +def _comma(arg): + if type(arg) is list: + arg = ','.join(arg) + return arg + + +def join(self, channels, keys=None): + channels = _comma(channels) + + if keys: + keys = _comma(keys) + self.write(('JOIN', channels, keys)) + else: + self.write(('JOIN', channels)) + +def nick(self, nickname): + self.write(('NICK', nickname)) + +def notice(self, msgtarget, message): + self.write(('NOTICE', msgtarget), message) + +def part(self, channels, message=None): + channels = _comma(channels) + self.write(('PART', channels), message) + +def pass_(self, password): + self.write(('PASS', password)) + +def ping(self, server1, server2=None): + self.write(('PING', server1), server2) + +def pong(self, server1, server2=None): + self.write(('PONG', server1), server2) + +def privmsg(self, msgtarget, message): + self.write(('PRIVMSG', msgtarget), message) + +def quit(self, message=None): + self.write(('QUIT'), message) + +def user(self, user, mode, realname): + self.write(('USER', user, mode, '_'), realname) + + +module_dict = sys.modules[__name__].__dict__ +command_filter = lambda k, v: callable(v) and not k.startswith('_') +commands = {k: v for k, v in module_dict.items() if command_filter(k, v)} diff --git a/test/test_irc.py b/test/test_irc.py index 8665d0c..7299b2c 100644 --- a/test/test_irc.py +++ b/test/test_irc.py @@ -42,7 +42,7 @@ class BotTest(unittest.TestCase): mock_write.assert_has_calls([ call(('NICK', self.nick)), - call(('USER', self.nick, '+iw', self.nick), self.name) + call(('USER', self.nick, '+iw', '_'), self.name) ]) @patch('irc.Bot.write') @@ -50,7 +50,7 @@ class BotTest(unittest.TestCase): self.bot.buffer = b"PING" self.bot.found_terminator() - mock_write.assert_called_once_with(('PONG', '')) + mock_write.assert_called_once_with(('PONG', ''), None) @patch('irc.Bot.push') def test_msg(self, mock_push): @@ -80,6 +80,6 @@ class BotTest(unittest.TestCase): @patch('irc.Bot.write') def test_notice(self, mock_write): notice = "This is a notice!" - self.bot.notice('jqh', notice) + self.bot.proto.notice('jqh', notice) mock_write.assert_called_once_with(('NOTICE', 'jqh'), notice) diff --git a/tools.py b/tools.py index d3a659e..6242168 100755 --- a/tools.py +++ b/tools.py @@ -8,6 +8,19 @@ http://inamidst.com/phenny/ """ +def decorate(obj, delegate): + class Decorator(object): + def __getattr__(self, attr): + if attr in delegate: + return delegate[attr] + + return getattr(obj, attr) + + def __setattr__(self, attr, value): + return setattr(obj, attr, value) + + return Decorator() + class GrumbleError(Exception): pass