From a50a5f019e1eb9d2fae08203e938d23eab196126 Mon Sep 17 00:00:00 2001 From: Yuuki Chan Date: Thu, 22 Feb 2024 18:41:50 +0900 Subject: [PATCH] Updated and fixed eltiebot.py, bot.py and sasl.py to use asyncio. --- elitebot.py | 3 +- src/bot.py | 134 ++++++++++++++++++++++++++++------------------------ src/sasl.py | 2 +- 3 files changed, 74 insertions(+), 65 deletions(-) diff --git a/elitebot.py b/elitebot.py index 4f2d53a..5bac82c 100755 --- a/elitebot.py +++ b/elitebot.py @@ -1,5 +1,6 @@ #!/usr/bin/env python3 +import asyncio import sys from src.bot import Bot @@ -22,7 +23,7 @@ def main(): try: print('EliteBot started successfully!') - bot.start() + asyncio.run(bot.start()) except Exception as e: print(f'Error starting EliteBot: {e}') sys.exit(1) diff --git a/src/bot.py b/src/bot.py index 635d45d..93cba71 100644 --- a/src/bot.py +++ b/src/bot.py @@ -1,13 +1,13 @@ #!/usr/bin/env python3 +import asyncio import importlib.util import inspect import json import os -import socket import ssl import sys -import time + import yaml from src.channel_manager import ChannelManager @@ -24,7 +24,8 @@ class Bot: self.channel_manager = ChannelManager() self.logger = Logger('logs/elitebot.log') self.connected = False - self.ircsock = None + self.reader = None + self.writer = None self.running = True self.plugins = [] self.load_plugins() @@ -32,7 +33,7 @@ class Bot: def validate_config(self, config): required_fields = [ ['Connection', 'Port'], - ['Connection' 'Hostname'], + ['Connection', 'Hostname'], ['Connection', 'Nick'], ['Connection', 'Ident'], ['Connection', 'Name'], @@ -92,11 +93,12 @@ class Bot: self.logger.error('Could not decode byte string with any known encoding') return bytes.decode('utf-8', 'ignore') - def ircsend(self, msg): + async def ircsend(self, msg): try: if msg != '': self.logger.info(f'Sending command: {msg}') - self.ircsock.send(bytes(f'{msg}\r\n', 'UTF-8')) + self.writer.write(bytes(f'{msg}\r\n', 'UTF-8')) + await self.writer.drain() except Exception as e: self.logger.error(f'Error sending IRC message: {e}') raise @@ -120,91 +122,97 @@ class Bot: args.append(' '.join(parts[trailing_arg_start:])[1:]) return source, command, args - def connect(self): + async def connect(self): try: - self.ircsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - + ssl_context = None if str(self.config['Connection'].get('Port'))[:1] == '+': - context = ssl.create_default_context() - self.ircsock = context.wrap_socket(self.ircsock, - server_hostname=self.config['Connection'].get('Hostname')) - port = int(self.config['Connection'].get('Port')[1:]) - else: - port = int(self.config['Connection'].get('Port')) + ssl_context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) # Corrected here - if 'BindHost' in self.config: - self.ircsock.bind((self.config['Connection'].get('BindHost'), 0)) + self.reader, self.writer = await asyncio.open_connection( + self.config['Connection'].get('Hostname'), + int(self.config['Connection'].get('Port')[1:]) if ssl_context else int( + self.config['Connection'].get('Port')), + ssl=ssl_context + ) - self.ircsock.connect_ex((self.config['Connection'].get('Hostname'), port)) - self.ircsend(f'NICK {self.config["Connection"].get("Nick")}') - self.ircsend(f'USER {self.config["Connection"].get("Ident")} * * :{self.config["Connection"].get("Name")}') + await self.ircsend(f'NICK {self.config["Connection"].get("Nick")}') + await self.ircsend( + f'USER {self.config["Connection"].get("Ident")} * * :{self.config["Connection"].get("Name")}') if self.config['SASL'].get('UseSASL'): - self.ircsend('CAP REQ :sasl') + await self.ircsend('CAP REQ :sasl') except Exception as e: self.logger.error(f'Error establishing connection: {e}') self.connected = False return - def start(self): + async def start(self): while True: if not self.connected: try: - self.connect() + await self.connect() self.connected = True except Exception as e: self.logger.error(f'Connection error: {e}') - time.sleep(60) + await asyncio.sleep(60) continue try: - recvText = self.ircsock.recv(2048) + recvText = await self.reader.read(2048) if not recvText: self.connected = False continue ircmsg = self.decode(recvText) - source, command, args = self.parse_message(ircmsg) - self.logger.debug(f'Received: source: {source} | command: {command} | args: {args}') - if command == 'PING': - nospoof = args[0][1:] if args[0].startswith(':') else args[0] - self.ircsend(f'PONG :{nospoof}') - continue + if '\r\n' in ircmsg: + messages = ircmsg.split('\r\n') + elif '\n' in ircmsg: + messages = ircmsg.split('\n') + else: + messages = [ircmsg] # If no newline characters, treat the whole message as a single message - if command == 'PRIVMSG': - channel, message = args[0], args[1] - source_nick = source.split('!')[0] - if message.startswith('&'): - cmd, *cmd_args = message[1:].split() - self.handle_command(source_nick, channel, cmd, cmd_args) - for plugin in self.plugins: - plugin.handle_message(source_nick, channel, message) + for message in messages: + if message: # Check if message is not empty + source, command, args = self.parse_message(message) + self.logger.debug(f'Received: source: {source} | command: {command} | args: {args}') - elif command == 'CAP' and args[1] == 'ACK' and 'sasl' in args[2]: - handle_sasl(self.config, self.ircsend) + match command: + case 'PING': + nospoof = args[0][1:] if args[0].startswith(':') else args[0] + await self.ircsend(f'PONG :{nospoof}') + continue + case 'PRIVMSG': + channel, message = args[0], args[1] + source_nick = source.split('!')[0] - elif command == 'AUTHENTICATE': - handle_authenticate(args, self.config, self.ircsend) - - elif command == '903': - handle_903(self.ircsend) - - if command == 'PRIVMSG' and args[1].startswith('\x01VERSION\x01'): - source_nick = source.split('!')[0] - self.ircsend(f'NOTICE {source_nick} :\x01VERSION EliteBot 0.1\x01') - - if command == '001': - for channel in self.channel_manager.get_channels(): - self.ircsend(f'JOIN {channel}') - - if command == 'INVITE': - channel = args[1] - self.ircsend(f'JOIN {channel}') - self.channel_manager.save_channel(channel) - - if command == 'VERSION': - self.ircsend('NOTICE', f'{source_nick} :I am a bot version 1.0.0') + if message.startswith('&'): + cmd, *cmd_args = message[1:].split() + self.handle_command(source_nick, channel, cmd, cmd_args) + elif args[1].startswith('\x01VERSION\x01'): + source_nick = source.split('!')[0] + await self.ircsend(f'NOTICE {source_nick} :\x01VERSION EliteBot 0.1\x01') + for plugin in self.plugins: + await plugin.handle_message(source_nick, channel, message) + case 'CAP': + if args[1] == 'ACK' and 'sasl' in args[2]: + handle_sasl(self.config, self.ircsend) + case 'AUTHENTICATE': + handle_authenticate(args, self.config, self.ircsend) + case 'INVITE': + channel = args[1] + await self.ircsend(f'JOIN {channel}') + self.channel_manager.save_channel(channel) + case 'VERSION': + await self.ircsend(f'NOTICE {source_nick} :I am a bot version 1.0.0') + case '001': + await self.ircsend('JOIN #YuukiTest') + # for channel in self.channel_manager.get_channels(): + # await self.ircsend(f'JOIN {channel}') + case '903': + await handle_903(self.ircsend) + case _: + continue except Exception as e: self.logger.error(f'General error: {e}') self.connected = False @@ -213,7 +221,7 @@ class Bot: if __name__ == '__main__': try: bot = Bot(sys.argv[1]) - bot.start() + asyncio.run(bot.start()) except KeyboardInterrupt: print('\nEliteBot has been stopped.') except Exception as e: diff --git a/src/sasl.py b/src/sasl.py index 513bf93..3581224 100644 --- a/src/sasl.py +++ b/src/sasl.py @@ -36,7 +36,7 @@ def handle_authenticate(args, config, ircsend): raise KeyError('SASLNICK and/or SASLPASS not found in config') -def handle_903(ircsend): +async def handle_903(ircsend): """ Handles the 903 command by sending a CAP END command.