diff --git a/.env b/.env new file mode 100644 index 0000000..c0337fd --- /dev/null +++ b/.env @@ -0,0 +1 @@ +USER_MAX_CHANNELS=50 diff --git a/paircd/client.py b/paircd/client.py index 9a92bc6..4ad6880 100644 --- a/paircd/client.py +++ b/paircd/client.py @@ -1,10 +1,12 @@ -from asyncio import StreamReader, StreamWriter, Queue +from asyncio import StreamReader, StreamWriter, TimeoutError, Queue +from asyncio.tasks import wait_for from dataclasses import dataclass, field from datetime import datetime from logging import log, INFO -from paircd.reply import RPL_CREATED, RPL_MYINFO, RPL_WELCOME, RPL_YOURHOST +from typing import Any, Set + +from paircd.reply import QUIT, RPL_CREATED, RPL_MYINFO, RPL_WELCOME, RPL_YOURHOST from paircd.message import Message -from typing import Set @dataclass @@ -12,6 +14,7 @@ class Client: hostname: str reader: StreamReader writer: StreamWriter + msg_queue: Queue = field(default_factory=Queue) nickname: str = "" @@ -19,6 +22,7 @@ class Client: realname: str = "" registered: bool = False away: str = "" + closed = False modes: Set[str] = field(default_factory=set) channels: Set[str] = field(default_factory=set) @@ -31,11 +35,43 @@ class Client: def log(self, msg: str, level: int = INFO) -> None: log(level, f"{self.hostname} ({self.id()}) {msg}") - async def write_forever(self) -> None: - while True: - msg = await self.msg_queue.get() - self.writer.write(msg) - await self.writer.drain() + async def write_until_closed(self, server: Any) -> None: + while not self.closed: + msg = None + try: + msg = await wait_for(self.msg_queue.get(), timeout=1.0) + except TimeoutError: + pass + + if msg is not None: + self.writer.write(msg) + try: + await self.writer.drain() + except ConnectionResetError: + await self.quit(server, "Connection reset by peer") + + async def quit(self, server: Any, msg: str) -> None: + if self.closed: + return + + quit_msg = QUIT(msg, prefix=self.id()) + for client in server.clients_by_nick.values(): + client.write_message(quit_msg) + for channel_name in self.channels: + channel = server.get_channel_by_name(channel_name, create=False) + if channel is None: + continue + channel.remove_client_by_nick(self.nickname) + server.remove_client_by_name(self.nickname) + await self.close() + + async def close(self) -> None: + if self.closed: + return + self.closed = True + if not self.writer.is_closing(): + self.writer.close() + await self.writer.wait_closed() def write_message(self, message: Message) -> None: self.msg_queue.put_nowait(message.encode()) diff --git a/paircd/handler/join.py b/paircd/handler/join.py index f3f9dab..398e1c5 100644 --- a/paircd/handler/join.py +++ b/paircd/handler/join.py @@ -1,5 +1,6 @@ from paircd.reply import ( ERR_NOSUCHCHANNEL, + ERR_TOOMANYCHANNELS, JOIN, RPL_ENDOFNAMES, RPL_NAMREPLY, @@ -11,6 +12,7 @@ from paircd.client import Client from paircd.command_handler import CommandHandler from paircd.message import Message from paircd.server import Server +from paircd.settings import USER_MAX_CHANNELS class JoinHandler(CommandHandler): @@ -20,6 +22,10 @@ class JoinHandler(CommandHandler): async def handle(self, server: Server, client: Client, msg: Message) -> None: channel_name = msg.args[0] + if USER_MAX_CHANNELS > 0 and len(client.channels) == USER_MAX_CHANNELS: + client.write_message(ERR_TOOMANYCHANNELS(client.nickname)) + return + channel = server.get_channel_by_name(channel_name) if channel is None: client.write_message(ERR_NOSUCHCHANNEL(client.nickname, channel_name)) diff --git a/paircd/handler/quit.py b/paircd/handler/quit.py new file mode 100644 index 0000000..5e426ce --- /dev/null +++ b/paircd/handler/quit.py @@ -0,0 +1,15 @@ +from paircd.client import Client +from paircd.command_handler import CommandHandler +from paircd.message import Message +from paircd.server import Server + + +class QuitHandler(CommandHandler): + def __init__(self) -> None: + super().__init__("QUIT") + + async def handle(self, server: Server, client: Client, msg: Message) -> None: + quit_msg = "" + if len(msg.args) == 1: + quit_msg = msg.args[0] + await client.quit(server, f"Quit: {quit_msg}") diff --git a/paircd/handlers.py b/paircd/handlers.py index c75194b..23ab279 100644 --- a/paircd/handlers.py +++ b/paircd/handlers.py @@ -15,6 +15,8 @@ from paircd.handler.nick import NickHandler from paircd.handler.notice import NoticeHandler from paircd.handler.ping import PingHandler from paircd.handler.privmsg import PrivmsgHandler + +from paircd.handler.quit import QuitHandler from paircd.handler.topic import TopicHandler from paircd.handler.user import UserHandler from paircd.handler.who import WhoHandler @@ -27,6 +29,7 @@ HANDLER_CLASSES = [ NoticeHandler, PingHandler, PrivmsgHandler, + QuitHandler, TopicHandler, UserHandler, WhoHandler, diff --git a/paircd/main.py b/paircd/main.py index 1ed8d4a..3e2b1eb 100644 --- a/paircd/main.py +++ b/paircd/main.py @@ -1,19 +1,31 @@ -from asyncio import create_task, run, start_server +from asyncio import create_task, run, start_server, IncompleteReadError from asyncio.streams import StreamReader, StreamWriter from logging import basicConfig, info, INFO +import logging from os import getenv +from sys import exc_info -from paircd.client import Client -from paircd.handlers import handle_cmd, register_cmd_handlers -from paircd.message import parse_message -from paircd.server import Server +from dotenv import load_dotenv + +load_dotenv("../.env") + +from paircd.client import Client # noqa: E402 +from paircd.handlers import handle_cmd, register_cmd_handlers # noqa: E402 +from paircd.message import parse_message # noqa: E402 +from paircd.server import Server # noqa: E402 basicConfig(format="%(asctime)s [%(levelname)s] - %(message)s", level=INFO) async def read_forever(server: Server, client: Client) -> None: - while True: - raw_msg = await client.reader.readuntil(b"\r\n") + while not client.closed: + try: + raw_msg = await client.reader.readuntil(b"\r\n") + except IncompleteReadError: + logging.warning("client connection closed", exc_info=exc_info()) + await client.quit(server, "Connection reset by peer") + break + msg = parse_message(raw_msg) await handle_cmd(server, client, msg) @@ -32,7 +44,7 @@ async def serve() -> None: writer=writer, ) create_task(read_forever(irc_server, client)) - create_task(client.write_forever()) + create_task(client.write_until_closed(irc_server)) server = await start_server( register_client, diff --git a/paircd/reply.py b/paircd/reply.py index 74e54bf..6af5df8 100644 --- a/paircd/reply.py +++ b/paircd/reply.py @@ -30,6 +30,7 @@ NICK = cmd_fn("NICK", "{0}") NOTICE = cmd_fn("NOTICE", "{0} :{1}") PONG = cmd_fn("PONG", ":{0}") PRIVMSG = cmd_fn("PRIVMSG", "{0} :{1}") +QUIT = cmd_fn("QUIT", ":{0}") TOPIC = cmd_fn("TOPIC", "{0} :{1}") # Error replies diff --git a/paircd/settings.py b/paircd/settings.py new file mode 100644 index 0000000..f648330 --- /dev/null +++ b/paircd/settings.py @@ -0,0 +1,3 @@ +import os + +USER_MAX_CHANNELS = int(os.getenv("USER_MAX_CHANNELS") or "0") diff --git a/poetry.lock b/poetry.lock index be40e66..4896637 100644 --- a/poetry.lock +++ b/poetry.lock @@ -188,6 +188,17 @@ wcwidth = "*" checkqa-mypy = ["mypy (==v0.761)"] testing = ["argcomplete", "hypothesis (>=3.56)", "mock", "nose", "requests", "xmlschema"] +[[package]] +name = "python-dotenv" +version = "0.18.0" +description = "Read key-value pairs from a .env file and set them as environment variables" +category = "main" +optional = false +python-versions = "*" + +[package.extras] +cli = ["click (>=5.0)"] + [[package]] name = "pyyaml" version = "5.4.1" @@ -249,7 +260,7 @@ python-versions = "*" [metadata] lock-version = "1.1" python-versions = "^3.8" -content-hash = "298291104b6892022cb47c34680f925bc85fff59e768a225230ef84444acb002" +content-hash = "16dd6949f4ee094d7e2d9ddf66b05ac484047dae65acbdd61c29e8a31862624c" [metadata.files] appdirs = [ @@ -345,6 +356,10 @@ pytest = [ {file = "pytest-5.4.3-py3-none-any.whl", hash = "sha256:5c0db86b698e8f170ba4582a492248919255fcd4c79b1ee64ace34301fb589a1"}, {file = "pytest-5.4.3.tar.gz", hash = "sha256:7979331bfcba207414f5e1263b5a0f8f521d0f457318836a7355531ed1a4c7d8"}, ] +python-dotenv = [ + {file = "python-dotenv-0.18.0.tar.gz", hash = "sha256:effaac3c1e58d89b3ccb4d04a40dc7ad6e0275fda25fd75ae9d323e2465e202d"}, + {file = "python_dotenv-0.18.0-py2.py3-none-any.whl", hash = "sha256:dd8fe852847f4fbfadabf6183ddd4c824a9651f02d51714fa075c95561959c7d"}, +] pyyaml = [ {file = "PyYAML-5.4.1-cp27-cp27m-macosx_10_9_x86_64.whl", hash = "sha256:3b2b1824fe7112845700f815ff6a489360226a5609b96ec2190a45e62a9fc922"}, {file = "PyYAML-5.4.1-cp27-cp27m-win32.whl", hash = "sha256:129def1b7c1bf22faffd67b8f3724645203b79d8f4cc81f674654d9902cb4393"}, diff --git a/pyproject.toml b/pyproject.toml index e707411..04e60a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,6 +6,7 @@ authors = ["Forest Belton <65484+forestbelton@users.noreply.github.com>"] [tool.poetry.dependencies] python = "^3.8" +python-dotenv = "^0.18.0" [tool.poetry.dev-dependencies] pytest = "^5.2"