paircd/paircd/message.py

from dataclasses import dataclass
from typing import List, Optional

MAX_MESSAGE_SIZE = 512

EXPECTED_ARG_COUNT = {
    "NICK": 1,
    "USER": 4,
    "JOIN": 1,
    "PRIVMSG": 2,
}


class IRCParsingError(Exception):
    message: str

    def __init__(self, message: Optional[str] = None) -> None:
        self.message = message


@dataclass
class IRCMessage:
    cmd: str
    args: List[str]
    prefix: str = ""

    def encode(self) -> bytes:
        prefix = self.prefix
        if prefix != "":
            prefix = f":{prefix} "

        # TODO: Raise exception if formatted message exceeds 512 bytes
        return f"{prefix}{self.cmd} {' '.join(self.args)}\r\n".encode("utf-8")

    @staticmethod
    def parse(raw: bytes) -> "IRCMessage":
        if len(raw) > MAX_MESSAGE_SIZE:
            raise IRCParsingError(
                f"Message is {len(raw)} bytes, larger than allowed {MAX_MESSAGE_SIZE}"
            )

        if not raw.endswith(b"\r\n"):
            raise IRCParsingError("Message does not terminate in CRLF")

        tokens: List[str] = []
        raw_tokens: List[str] = raw.decode("utf-8").split(" ")

        for i, token in enumerate(raw_tokens):
            if token.startswith(":"):
                trailing = token[1:] + " " + " ".join(raw_tokens[i + 1 :])
                tokens.append(trailing)
                break
            tokens.append(token)

        if len(tokens) == 0:
            raise IRCParsingError("Message has no command")

        cmd = tokens[0]
        if cmd in EXPECTED_ARG_COUNT and EXPECTED_ARG_COUNT[cmd] != len(tokens) - 1:
            raise IRCParsingError(
                f"{cmd} had {len(tokens)-1} arguments, expected {EXPECTED_ARG_COUNT[cmd]}"
            )

        tokens[-1] = tokens[-1].strip()
        return IRCMessage(cmd=tokens[0], args=tokens[1:])