From 68426970644d7602aa0ee400c3cb29658c4a75de Mon Sep 17 00:00:00 2001 From: luk3yx Date: Thu, 12 May 2022 16:40:31 +1200 Subject: Add CAP negotiation to idc_irc_proxy.py --- idc_irc_proxy.py | 115 ++++++++++++++++++++++++++++++++----------------------- 1 file changed, 68 insertions(+), 47 deletions(-) diff --git a/idc_irc_proxy.py b/idc_irc_proxy.py index 83eee0a..c02573b 100755 --- a/idc_irc_proxy.py +++ b/idc_irc_proxy.py @@ -7,7 +7,8 @@ import miniirc, miniirc_idc, os, socket, threading, traceback from concurrent.futures import ThreadPoolExecutor -from miniirc_extras.utils import ircv2_message_unparser, ircv3_message_parser +from miniirc_extras.utils import (ircv2_message_unparser, ircv3_message_parser, + ircv3_message_unparser) # A single network class Proxy: @@ -16,22 +17,41 @@ class Proxy: encoding = 'utf-8' _001 = False - _main_lock = False block_incoming = frozenset(('PING', 'CAP', 'AUTHENTICATE')) block_outgoing = frozenset(('CAP',)) + _advertised_caps = frozenset(( + 'server-time', 'message-tags', 'account-tag', 'echo-message' + )) + + # Create the IRC object + def __init__(self, conn, *args, **kwargs): + self.sock = conn + self.irc = self.IRC(*args, auto_connect=False, + executor=ThreadPoolExecutor(1), **kwargs) + self._irc_msg_unparser = ircv2_message_unparser + + # Add the IRC handler + self.irc.CmdHandler(ircv3=True, colon=False)(self._miniirc_handler) + + # Start the main loop + self.thread = threading.Thread(target=self._init_thread) + self.thread.start() # Send messages def send(self, cmd, hostmask, tags, args): - raw = ircv2_message_unparser(cmd, hostmask or (cmd, cmd, cmd), tags, - args, colon=False, encoding=self.encoding) - self.sock.sendall(raw[:510] + b'\r\n') + raw = self._irc_msg_unparser( + cmd, hostmask or (cmd, cmd, cmd), tags, args, colon=False, + encoding=self.encoding + ) + self.sock.sendall(raw + b'\r\n') # Receive messages def recv(self): while True: while b'\n' not in self._buffer: buf = self.sock.recv(4096) - assert buf, 'The socket has been closed!' + if not buf: + raise BrokenPipeError self._buffer += buf.replace(b'\r', b'\n') msg, self._buffer = self._buffer.split(b'\n', 1) @@ -56,7 +76,7 @@ class Proxy: self._sendcmd(*self._sendq.pop(0)) # Start the main loop - self._main() + threading.Thread(target=self._main).start() elif cmd == 'ERROR': self.send('PING', None, {}, [':ERROR']) elif cmd == 'PONG' and args and args[-1] == 'miniirc-ping': @@ -65,7 +85,7 @@ class Proxy: # Send the command to the client try: self.send(cmd, hostmask, tags, args) - except Exception as e: + except Exception: traceback.print_exc() self.irc.disconnect('Connection closed.', auto_reconnect=False) @@ -74,14 +94,48 @@ class Proxy: self._sendq = [] nick = None user = None + cap_negotiation = False # Wait for NICK and USER to be sent - while not nick or not user: - tags, cmd, args = self.recv() + while cap_negotiation or not nick or not user: + try: + tags, cmd, args = self.recv() + except BrokenPipeError: + return + if cmd == 'NICK' and len(args) == 1: nick = args[0] elif cmd == 'USER' and len(args) > 1: user = args + elif cmd == 'CAP' and args: + subcmd = args[0].upper() + if subcmd == 'LS': + cap_negotiation = True + self.send('CAP', None, {}, + ('*', 'LS', ' '.join(self._advertised_caps))) + elif subcmd == 'REQ': + cap_negotiation = True + requested_caps = frozenset(args[-1].split(' ')) + if requested_caps - self._advertised_caps: + self.send('CAP', None, {}, ('*', 'NAK', args[-1])) + continue + + if 'message-tags' in requested_caps: + self._irc_msg_unparser = ircv3_message_unparser + + if 'echo-message' in requested_caps: + self.irc.ircv3_caps.add('echo-message') + + self.send('CAP', None, {}, ('*', 'ACK', args[-1])) + elif subcmd == 'END': + cap_negotiation = False + else: + self.send('410', None, {}, + ('*', subcmd, 'Invalid CAP subcommand')) + elif cmd == 'QUIT': + self.sock.shutdown(socket.SHUT_RDWR) + self.sock.close() + return else: self._sendq.append((tags, cmd, args)) @@ -102,50 +156,16 @@ class Proxy: # The more permanent main loop def _main(self, single_thread=False): - if not single_thread: - if self._main_lock and self._main_lock.is_alive(): - return self._main_lock - - t = threading.Thread(target=self._main, args=(True,)) - t.start() - return t - - # Clear the RecvQ - if self._recvq: - while len(self._recvq) > 0: - self.send(*self._recvq.pop(0)) - self._recvq = None - # Send everything to IRC while True: try: tags, cmd, args = self.recv() - except Exception as e: - print(repr(e)) + except Exception: + traceback.print_exc() return self.irc.disconnect() self._sendcmd(tags, cmd, args) - # Generic init function - def _init(self, conn, irc): - self.sock = conn - self.irc = irc - - # Add the IRC handler - self._recvq = [] - self.irc.CmdHandler(ircv3=True, colon=False)(self._miniirc_handler) - - # Start the main loop - self.thread = threading.Thread(target=self._init_thread) - self.thread.start() - - # Create the IRC object - def __init__(self, conn, *args, bad_cmds=None, **kwargs): - if bad_cmds is not None: - self.bad_cmds = bad_cmds - self._init(conn, self.IRC(*args, auto_connect=False, - executor=ThreadPoolExecutor(1), **kwargs)) - # The proxy class class Server: Proxy = Proxy @@ -173,10 +193,11 @@ def main(): parser.add_argument('local_port', type=int) parser.add_argument('username') parser.add_argument('password') + parser.add_argument('--debug', '-v', action='store_true') args = parser.parse_args() Server('andrewyu.org', 6835, '', ssl=True, persist=False, - local_addr=('127.0.0.1', args.local_port), + local_addr=('127.0.0.1', args.local_port), debug=args.debug, ns_identity=(args.username, args.password)).main() if __name__ == '__main__': -- cgit v1.2.3