ircrobots/ircrobots/ircv3.py

from typing    import Iterable, List, Optional
from irctokens import build

from .contexts import  ServerContext
from .matching import  (Response, Numerics, ResponseOr, ParamAny, ParamNot,
    ParamLiteral)
from .interface import ICapability

class Capability(ICapability):
    def __init__(self,
            ratified_name: Optional[str],
            draft_name:    Optional[str]=None,
            alias:         Optional[str]=None,
            depends_on:    List[str]=[]):
        self.name  = ratified_name
        self.draft = draft_name
        self.alias = alias or ratified_name
        self.depends_on = depends_on.copy()

        self._caps = set((ratified_name, draft_name))

    def available(self, capabilities: Iterable[str]
            ) -> Optional[str]:
        match = list(set(capabilities)&self._caps)
        return match[0] if match else None

    def match(self, capability: str) -> Optional[str]:
        cap = list(set([capability])&self._caps)
        return cap[0] if cap else None

    def copy(self):
        return Capability(
            self.name,
            self.draft,
            alias=self.alias,
            depends_on=self.depends_on[:])

CAP_SASL = Capability("sasl")
CAPS: List[ICapability] = [
    Capability("multi-prefix"),
    Capability("chghost"),
    Capability("away-notify"),
    Capability("userhost-in-names"),

    Capability("invite-notify"),
    Capability("account-tag"),
    Capability("account-notify"),
    Capability("extended-join"),

    Capability("message-tags", "draft/message-tags-0.2"),
    Capability("cap-notify"),
    Capability("batch"),

    Capability(None, "draft/rename", alias="rename"),
    Capability("setname", "draft/setname")
]

class CAPContext(ServerContext):
    async def handshake(self) -> bool:
        # improve this by being able to wait_for Emit objects
        line = await self.server.wait_for(ResponseOr(
            Response(
                "CAP",
                [ParamAny(), ParamLiteral("LS"), ParamNot(ParamLiteral("*"))]
            ),
            Numerics(["RPL_WELCOME"])
        ))

        if line.command == "CAP":
            caps = self.server.collect_caps()
            if caps:
                await self.server.send(
                    build("CAP", ["REQ", " ".join(caps)]))

                while caps:
                    line = await self.server.wait_for(ResponseOr(
                        Response("CAP", [ParamAny(), ParamLiteral("ACK")]),
                        Response("CAP", [ParamAny(), ParamLiteral("NAK")])
                    ))

                    current_caps = line.params[2].split(" ")
                    for cap in current_caps:
                        if cap in caps:
                            caps.remove(cap)

            if (self.server.cap_agreed(CAP_SASL) and
                    not self.server.params.sasl is None):
                await self.server.sasl_auth(self.server.params.sasl)

            await self.server.send(build("CAP", ["END"]))
            return True
        else:
            return False