Merge pull request #173 from TotallyNotRobots/channel-keys

Add official support for channel keys

Author: linuxdaemon

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Add Python 3.8 to testing matrix
 - Add support for channel keys (#95)
+- Officially support channel keys across the whole bot
 ### Changed
 - Refactor tests to remove dependency on mock library
 - Change link_announcer.py to only warn on connection errors

cloudbot/clients/irc.py

@@ -146,6 +146,35 @@ def __init__(self, bot, _type, name, nick, *, channels=None, config=None):
 
         self._connecting = False
 
+        self._channel_keys = {}
+
+    def set_channel_key(self, channel: str, key: str, *, override: bool = True) -> None:
+        if override or channel not in self._channel_keys:
+            self._channel_keys[channel] = key
+
+    def clear_channel_keys(self) -> None:
+        self._channel_keys.clear()
+
+    def clear_channel_key(self, channel: str) -> bool:
+        if channel in self._channel_keys:
+            del self._channel_keys[channel]
+            return True
+
+        return False
+
+    def get_channel_key(
+        self, channel: str, default: Optional[str] = None, *, set_key: bool = True
+    ) -> Optional[str]:
+        if channel in self._channel_keys:
+            key = self._channel_keys[channel]
+            if key is not None:
+                return key
+
+        if set_key:
+            self._channel_keys[channel] = default
+
+        return default
+
     def make_ssl_context(self, conn_config):
         if self.use_ssl:
             ssl_context = ssl.create_default_context()
@@ -304,6 +333,7 @@ def set_nick(self, nick):
         self.cmd("NICK", nick)
 
     def join(self, channel, key=None):
+        key = self.get_channel_key(channel, key)
         if key:
             self.cmd("JOIN", channel, key)
         else:

cloudbot/util/irc.py

@@ -0,0 +1,86 @@
+import logging
+from enum import Enum
+from typing import List, Mapping, Optional
+
+import attr
+
+logger = logging.getLogger(__name__)
+
+
+class ModeType(Enum):
+    A = "A"
+    B = "B"
+    C = "C"
+    D = "D"
+    Status = 1
+
+
+PARAM_MODE_TYPES = (ModeType.A, ModeType.B, ModeType.Status)
+
+
+@attr.s(hash=True)
+class ChannelMode:
+    """
+    An IRC channel mode
+    """
+
+    character = attr.ib(type=str)
+    type = attr.ib(type=ModeType)
+
+    def has_param(self, adding: bool) -> bool:
+        return self.type in PARAM_MODE_TYPES or (self.type is ModeType.C and adding)
+
+
+@attr.s(hash=True)
+class ModeChange:
+    """
+    Represents a single change of a mode
+    """
+
+    char = attr.ib(type=str)
+    adding = attr.ib(type=bool)
+    param = attr.ib(type=Optional[str])
+    info = attr.ib(type=ChannelMode)
+
+    @property
+    def is_status(self):
+        return self.info.type is ModeType.Status
+
+
+@attr.s(hash=True)
+class StatusMode(ChannelMode):
+    """
+    An IRC status mode
+    """
+
+    prefix = attr.ib(type=str)
+    level = attr.ib(type=int)
+
+    @classmethod
+    def make(cls, prefix: str, char: str, level: int) -> "StatusMode":
+        return cls(prefix=prefix, level=level, character=char, type=ModeType.Status)
+
+
+def parse_mode_string(
+    modes: str, params: List[str], server_modes: Mapping[str, ChannelMode]
+) -> List[ModeChange]:
+    new_modes = []
+    params = params.copy()
+    adding = True
+    for c in modes:
+        if c == "+":
+            adding = True
+        elif c == "-":
+            adding = False
+        else:
+            mode_info = server_modes.get(c)
+            if mode_info and mode_info.has_param(adding):
+                param = params.pop(0)
+            else:
+                param = None
+
+            new_modes.append(
+                ModeChange(char=c, adding=adding, param=param, info=mode_info)
+            )
+
+    return new_modes

plugins/core/chan_key_db.py

@@ -0,0 +1,122 @@
+"""
+Store and retrieve channel keys in a database table
+
+Author:
+- linuxdaemon
+"""
+from itertools import zip_longest
+from typing import Any, Dict, List, Optional
+
+from irclib.parser import Message
+from sqlalchemy import Column, PrimaryKeyConstraint, String, Table, and_, select
+from sqlalchemy.orm import Session
+from sqlalchemy.sql.elements import BooleanClauseList, ClauseElement
+
+from cloudbot import hook
+from cloudbot.client import Client
+from cloudbot.clients.irc import IrcClient
+from cloudbot.util import database
+from cloudbot.util.irc import parse_mode_string
+from plugins.core.server_info import get_channel_modes, get_server_info
+
+table = Table(
+    "channel_keys",
+    database.metadata,
+    Column("conn", String),
+    Column("chan", String),
+    Column("key", String),
+    PrimaryKeyConstraint("conn", "chan"),
+)
+
+
+@hook.connect(clients=["irc"])
+def load_keys(conn: IrcClient, db) -> None:
+    """
+    Load channel keys to the client
+    """
+    query = select([table.c.chan, table.c.key], table.c.conn == conn.name.lower())
+    conn.clear_channel_keys()
+    for row in db.execute(query):
+        conn.set_channel_key(row["chan"], row["key"])
+
+
+@hook.irc_raw("MODE")
+def handle_modes(irc_paramlist: List[str], conn: IrcClient, db, chan: str) -> None:
+    """
+    Handle mode changes
+    """
+    modes = irc_paramlist[1]
+    mode_params = list(irc_paramlist[2:])
+    serv_info = get_server_info(conn)
+    mode_changes = parse_mode_string(modes, mode_params, get_channel_modes(serv_info))
+    updated = False
+    for change in mode_changes:
+        if change.char == "k":
+            updated = True
+            if change.adding:
+                set_key(db, conn, chan, change.param)
+            else:
+                clear_key(db, conn, chan)
+
+    if updated:
+        load_keys(conn, db)
+
+
+def insert_or_update(
+    db: Session, tbl: Table, data: Dict[str, Any], query: ClauseElement
+) -> None:
+    """
+    Insert a new row or update an existing matching row
+    """
+    result = db.execute(tbl.update().where(query).values(**data))
+    if not result.rowcount:
+        db.execute(tbl.insert().values(**data))
+
+    db.commit()
+
+
+def make_clause(conn: Client, chan: str) -> BooleanClauseList:
+    """
+    Generate a WHERE clause to match keys for this conn+channel
+    """
+    return and_(table.c.conn == conn.name.lower(), table.c.chan == chan.lower(),)
+
+
+def clear_key(db: Session, conn, chan: str) -> None:
+    """
+    Remove a channel's key from the DB
+    """
+    db.execute(table.delete().where(make_clause(conn, chan)))
+
+
+def set_key(db: Session, conn: IrcClient, chan: str, key: Optional[str]) -> None:
+    """
+    Set the key for a channel
+    """
+    insert_or_update(
+        db,
+        table,
+        {"conn": conn.name.lower(), "chan": chan.lower(), "key": key or None},
+        make_clause(conn, chan),
+    )
+    conn.set_channel_key(chan, key)
+
+
+@hook.irc_out()
+def check_send_key(conn: IrcClient, parsed_line: Message, db: Session) -> Message:
+    """
+    Parse outgoing JOIN messages and store used channel keys
+    """
+    if parsed_line.command == "JOIN":
+        if len(parsed_line.parameters) > 1:
+            keys = parsed_line.parameters[1]
+        else:
+            keys = ""
+
+        for chan, key in zip_longest(
+            *map(lambda s: s.split(","), (parsed_line.parameters[0], keys))
+        ):
+            if key:
+                set_key(db, conn, chan, key)
+
+    return parsed_line