From a4f0d1bf287762f6097c570ed11cc92ab5ef7131 Mon Sep 17 00:00:00 2001
From: jesopo <github@lolnerd.net>
Date: Fri, 7 Sep 2018 16:34:51 +0100
Subject: [PATCH] Support IRCv3's tls/STARTTLS

---
 IRCServer.py        | 15 ++++++++++-----
 modules/sasl.py     |  6 ------
 modules/starttls.py | 26 ++++++++++++++++++++++++++
 3 files changed, 36 insertions(+), 11 deletions(-)
 create mode 100644 modules/starttls.py

diff --git a/IRCServer.py b/IRCServer.py
index 30e29580..8059a9a5 100644
--- a/IRCServer.py
+++ b/IRCServer.py
@@ -65,11 +65,7 @@ class Server(object):
         self.socket.settimeout(5.0)
 
         if self.tls:
-            context = ssl.SSLContext(OUR_TLS_PROTOCOL)
-            context.options |= ssl.OP_NO_SSLv2
-            context.options |= ssl.OP_NO_SSLv3
-            context.options |= ssl.OP_NO_TLSv1
-            self.socket = context.wrap_socket(self.socket)
+            self.tls_wrap()
         self.cached_fileno = self.socket.fileno()
         self.events.on("timer").on("rejoin").hook(self.try_rejoin)
 
@@ -82,6 +78,13 @@ class Server(object):
         fileno = self.socket.fileno()
         return self.cached_fileno if fileno == -1 else fileno
 
+    def tls_wrap(self):
+        context = ssl.SSLContext(OUR_TLS_PROTOCOL)
+        context.options |= ssl.OP_NO_SSLv2
+        context.options |= ssl.OP_NO_SSLv3
+        context.options |= ssl.OP_NO_TLSv1
+        self.socket = context.wrap_socket(self.socket)
+
     def connect(self):
         self.socket.connect((self.target_hostname, self.port))
         self.send_capibility_ls()
@@ -297,6 +300,8 @@ class Server(object):
         self.send("CAP END")
     def send_authenticate(self, text):
         self.send("AUTHENTICATE %s" % text)
+    def send_starttls(self):
+        self.send("STARTTLS")
 
     def waiting_for_capabilities(self):
         return bool(len(self._capabilities_waiting))
diff --git a/modules/sasl.py b/modules/sasl.py
index 51e48587..32e20137 100644
--- a/modules/sasl.py
+++ b/modules/sasl.py
@@ -3,17 +3,11 @@ import base64
 class Module(object):
     def __init__(self, bot, events, exports):
         self.bot = bot
-        events.on("preprocess.connect").hook(self.preprocess_connect)
         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.903").hook(self.sasl_success)
 
-    def preprocess_connect(self, event):
-        sasl = event["server"].get_setting("sasl")
-        if sasl:
-            event["server"].send_capability_request("sasl")
-
     def on_cap(self, event):
         has_sasl = "sasl" in event["capabilities"]
         has_mechanisms = has_sasl and not event["capabilities"]["sasl"
diff --git a/modules/starttls.py b/modules/starttls.py
new file mode 100644
index 00000000..e3e69341
--- /dev/null
+++ b/modules/starttls.py
@@ -0,0 +1,26 @@
+import base64
+
+class Module(object):
+    def __init__(self, bot, events, exports):
+        self.bot = bot
+        events.on("received.cap.ls").hook(self.on_cap)
+        events.on("received.cap.ack").hook(self.on_cap_ack)
+
+        events.on("received.numeric.670").hook(self.starttls_success)
+        events.on("received.numeric.691").hook(self.starttls_failed)
+
+    def on_cap(self, event):
+        if "tls" in event["capabilities"].keys() and not event["server"].tls:
+            event["server"].queue_capability("tls")
+
+    def on_cap_ack(self, event):
+        if "tls" in event["capabilities"].keys():
+            event["server"].send_starttls()
+            event["server"].wait_for_capability("tls")
+
+    def starttls_success(self, event):
+        event["server"].wrap_tls()
+        event["server"].capability_done("tls")
+    def starttls_failed(self, event):
+        event["server"].capability_done("tls")
+