Make CAP functionality more part of the framework and support message-tags and

multi-prefix

3 files changed, 59 insertions, 23 deletions

diff --git a/IRCLineHandler.py b/IRCLineHandler.py
index 6129b42e..d8958795 100644
--- a/IRCLineHandler.py
+++ b/IRCLineHandler.py
@@ -43,8 +43,16 @@ class LineHandler(object):
 
     def handle(self, server, line):
         original_line = line
+        tags = {}
         prefix = None
         command = None
+
+        if line[0] == "@":
+            tags_prefix, line = line[1:].split(" ", 1)
+            for tag in tags_prefix.split(";"):
+                tag_split = tag.split("=", 1)
+                tags[tag[0]] = tag[1:]
+
         if line[0] == ":":
             prefix, command, line = line[1:].split(" ", 2)
         else:
@@ -66,16 +74,16 @@ class LineHandler(object):
 
         #server, prefix, command, args, arbitrary
         self.events.on("raw").on(command).call(server=server,
-            prefix=prefix, args=args, arbitrary=arbitrary)
+            prefix=prefix, args=args, arbitrary=arbitrary, tags=tags)
         if default_event or not hooks:
             if command.isdigit():
                 self.events.on("received").on("numeric").on(command
-                    ).call(line=original_line, server=server,
+                    ).call(line=original_line, server=server, tags=tags,
                     line_split=original_line.split(" "), number=command)
             else:
                 self.events.on("received").on(command).call(
                     line=original_line, line_split=original_line.split(" "),
-                    command=command, server=server)
+                    command=command, server=server, tags=tags)
 
     # ping from the server
     def ping(self, event):
@@ -243,9 +251,22 @@ class LineHandler(object):
     # the server is telling us about its capabilities!
     def cap(self, event):
         capabilities = (event["arbitrary"] or "").split(" ")
-        self.events.on("received").on("cap").call(
-            subcommand=event["args"][1], capabilities=capabilities,
-            server=event["server"])
+        subcommand = event["args"][1].lower()
+
+        if subcommand == "ls":
+            if "message-tags" in capabilities:
+                event["server"].queue_capability("message-tags")
+            if "multi-prefix" in capabilities:
+                event["server"].queue_capability("multi-prefix")
+
+        self.events.on("received").on("cap").on(subcommand).call(
+            capabilities=capabilities, server=event["server"])
+
+        if subcommand == "ls":
+            if event["server"].has_capability_queue():
+                event["server"].send_capability_queue()
+            else:
+                event["server"].send_capability_end()
 
     # the server is asking for authentication
     def authenticate(self, event):
diff --git a/IRCServer.py b/IRCServer.py
index 34358e44..9cb5e18f 100644
--- a/IRCServer.py
+++ b/IRCServer.py
@@ -26,6 +26,8 @@ class Server(object):
         self.original_username = username or nickname
         self.original_realname = realname or nickname
         self.name = None
+        self._capability_queue = set([])
+        self._capabilities_waiting = set([])
 
         self.write_buffer = b""
         self.buffered_lines = []
@@ -79,7 +81,7 @@ class Server(object):
 
     def connect(self):
         self.socket.connect((self.target_hostname, self.port))
-        self.events.on("preprocess.connect").call(server=self)
+        self.send_capibility_ls()
 
         if self.password:
             self.send_pass(self.password)
@@ -273,13 +275,31 @@ class Server(object):
     def send_nick(self, nickname):
         self.send("NICK %s" % nickname)
 
-    def send_capability_request(self, capname):
-        self.send("CAP REQ :%s" % capname)
+    def send_capibility_ls(self):
+        self.send("CAP LS")
+    def queue_capability(self, capability):
+        self._capability_queue.add(capability)
+    def send_capability_queue(self):
+        if self.has_capability_queue():
+            capabilities = " ".join(self._capability_queue)
+            self._capability_queue.clear()
+            self.send_capability_request(capabilities)
+    def has_capability_queue(self):
+        return bool(len(self._capability_queue))
+    def send_capability_request(self, capability):
+        self.send("CAP REQ :%s" % capability)
     def send_capability_end(self):
         self.send("CAP END")
     def send_authenticate(self, text):
         self.send("AUTHENTICATE %s" % text)
 
+    def wait_for_capability(self, capability):
+        self._capabilities_waiting.add(capability)
+    def capability_done(self, capability):
+        self._capabilities_waiting.remove(capability)
+        if not self._capabilities_waiting:
+            self.send_capability_end()
+
     def send_pass(self, password):
         self.send("PASS %s" % password)
 
diff --git a/modules/sasl.py b/modules/sasl.py
index 15a7523a..a5092615 100644
--- a/modules/sasl.py
+++ b/modules/sasl.py
@@ -4,10 +4,9 @@ class Module(object):
     def __init__(self, bot, events, exports):
         self.bot = bot
         events.on("preprocess.connect").hook(self.preprocess_connect)
-        events.on("received.cap").hook(self.on_cap)
+        events.on("received.cap.ls").hook(self.on_cap)
+        events.on("received.cap.ack").hook(self.on_cap_ack)
         events.on("received.authenticate").hook(self.on_authenticate)
-        events.on("received.numeric").on(
-            "902", "903", "904", "905", "906", "907", "908").hook(self.on_90x)
 
     def preprocess_connect(self, event):
         sasl = event["server"].get_setting("sasl")
@@ -15,13 +14,12 @@ class Module(object):
             event["server"].send_capability_request("sasl")
 
     def on_cap(self, event):
-        if event["subcommand"] == "NAK":
-            event["server"].send_capability_end()
-        elif event["subcommand"] == "ACK":
-            if not "sasl" in event["capabilities"]:
-                event["server"].send_capability_end()
-            else:
-                event["server"].send_authenticate("PLAIN")
+        if "sasl" in event["capabilities"]:
+            event["server"].queue_capability("sasl")
+    def on_cap_ack(self, event):
+        if "sasl" in event["capabilities"]:
+            event["server"].send_authenticate("PLAIN")
+            event["server"].wait_for_capability("sasl")
 
     def on_authenticate(self, event):
         if event["message"] != "+":
@@ -33,7 +31,4 @@ class Module(object):
             auth_text = base64.b64encode(auth_text.encode("utf8"))
             auth_text = auth_text.decode("utf8")
             event["server"].send_authenticate(auth_text)
-
-    def on_90x(self, event):
-        event["server"].send_capability_end()
-
+        event["server"].capability_done("sasl")