diff --git a/README.md b/README.md new file mode 100644 index 0000000..0a16bf5 --- /dev/null +++ b/README.md @@ -0,0 +1,10 @@ +# ircrobots + +## rationale +I wanted a very-bare-bones IRC bot framework that deals with most of the +concerns one would deal with in scheduling and awaiting async stuff, e.g. +creating and awaiting a new task for each server while dynamically being able +to add/remove servers. + +## usage +see [examples/](examples/) for some usage demonstration. diff --git a/examples/simple.py b/examples/simple.py new file mode 100644 index 0000000..f04f83e --- /dev/null +++ b/examples/simple.py @@ -0,0 +1,26 @@ +import asyncio + +from irctokens import build, Line + +from ircrobots.bot import Bot as BaseBot +from ircrobots.server import ConnectionParams, Server + +SERVERS = [ + ("freenode", "chat.freenode.net"), + ("tilde", "ctrl-c.tilde.chat") +] + +class Bot(BaseBot): + async def line_read(self, server: Server, line: Line): + if line.command == "001": + print(f"connected to {server.isupport.network}") + await server.send(build("JOIN", ["#testchannel"])) + +async def main(): + bot = Bot() + for name, host in SERVERS: + params = ConnectionParams("BitBotNewTest", host, 6697, True) + await bot.add_server(name, params) + await bot.run() +if __name__ == "__main__": + asyncio.run(main()) diff --git a/ircrobots/__init__.py b/ircrobots/__init__.py new file mode 100644 index 0000000..8b13789 --- /dev/null +++ b/ircrobots/__init__.py @@ -0,0 +1 @@ + diff --git a/ircrobots/bot.py b/ircrobots/bot.py new file mode 100644 index 0000000..bef2e11 --- /dev/null +++ b/ircrobots/bot.py @@ -0,0 +1,62 @@ +import asyncio, inspect +import anyio +from queue import Queue +from typing import Any, Awaitable, Callable, cast, Dict, List, Tuple +from irctokens import Line + +from .server import ConnectionParams, Server + +RECONNECT_DELAY = 10.0 # ten seconds reconnect + +class Bot(object): + def __init__(self): + self.servers: Dict[str, Server] = {} + self._server_queue: asyncio.Queue[Server] = asyncio.Queue() + + # methods designed to be overridden + def create_server(self, name: str): + return Server(name) + async def disconnected(self, server: Server): + await asyncio.sleep(RECONNECT_DELAY) + await self.add_server(server.name, server.params) + async def line_read(self, server: Server, line: Line): + pass + async def line_send(self, server: Server, line: Line): + pass + + async def add_server(self, name: str, params: ConnectionParams) -> Server: + server = self.create_server(name) + self.servers[name] = server + await server.connect(params) + await self._server_queue.put(server) + return server + + async def _run_server(self, server: Server): + async with anyio.create_task_group() as tg: + async def _read(): + while not tg.cancel_scope.cancel_called: + lines = await server._read_lines() + for line in lines: + await self.line_read(server, line) + await tg.cancel_scope.cancel() + + async def _write(): + try: + while not tg.cancel_scope.cancel_called: + lines = await server._write_lines() + for line in lines: + await self.line_send(server, line) + except Exception as e: + print(e) + await tg.cancel_scope.cancel() + + await tg.spawn(_read) + await tg.spawn(_write) + del self.servers[server.name] + await self.disconnected(server) + + async def run(self): + async with anyio.create_task_group() as tg: + while not tg.cancel_scope.cancel_called: + server = await self._server_queue.get() + await tg.spawn(self._run_server, server) diff --git a/ircrobots/server.py b/ircrobots/server.py new file mode 100644 index 0000000..5018a37 --- /dev/null +++ b/ircrobots/server.py @@ -0,0 +1,94 @@ +import asyncio, ssl +from queue import Queue +from typing import Callable, Dict, List, Optional, Tuple +from enum import Enum +from dataclasses import dataclass + +from asyncio_throttle import Throttler +from ircstates import Server as BaseServer +from irctokens import build, Line, tokenise + +sc = ssl.create_default_context(ssl.Purpose.SERVER_AUTH) + +THROTTLE_RATE = 4 # lines +THROTTLE_TIME = 2 # seconds + +@dataclass +class ConnectionParams(object): + nickname: str + host: str + port: int + ssl: bool + + username: Optional[str] = None + realname: Optional[str] = None + bindhost: Optional[str] = None + +class SendPriority(Enum): + HIGH = 0 + MEDIUM = 10 + LOW = 20 + + DEFAULT = MEDIUM + +class Server(BaseServer): + _reader: asyncio.StreamReader + _writer: asyncio.StreamWriter + params: ConnectionParams + + def __init__(self, name: str): + super().__init__(name) + self.throttle = Throttler( + rate_limit=THROTTLE_RATE, period=THROTTLE_TIME) + self._write_queue: asyncio.PriorityQueue[Tuple[int, Line]] = asyncio.PriorityQueue() + + async def send_raw(self, line: str, priority=SendPriority.DEFAULT): + await self.send(tokenise(line), priority) + async def send(self, line: Line, priority=SendPriority.DEFAULT): + await self._write_queue.put((priority, line)) + + def set_throttle(self, rate: int, time: float): + self.throttle.rate_limit = rate + self.throttle.period = time + + async def connect(self, params: ConnectionParams): + cur_ssl = sc if params.ssl else None + reader, writer = await asyncio.open_connection( + params.host, params.port, ssl=cur_ssl) + self._reader = reader + self._writer = writer + + nickname = params.nickname + username = params.username or nickname + realname = params.realname or nickname + + await self.send(build("NICK", [nickname])) + await self.send(build("USER", [username, "0", "*", realname])) + + self.params = params + + async def line_received(self, line: Line): + pass + async def _read_lines(self) -> List[Line]: + data = await self._reader.read(1024) + lines = self.recv(data) + for line in lines: + print(f"{self.name}< {line.format()}") + await self.line_received(line) + return lines + + async def line_written(self, line: Line): + pass + async def _write_lines(self) -> List[Line]: + lines: List[Line] = [] + + while (not lines or + (len(lines) < 5 and self._write_queue.qsize() > 0)): + prio, line = await self._write_queue.get() + lines.append(line) + + for line in lines: + async with self.throttle: + self._writer.write(f"{line.format()}\r\n".encode("utf8")) + await self._writer.drain() + return lines diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..8187ad2 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +ircstates ==0.7.0