Abstract ssl socket wrapping away so we can reuse it

3 files changed, 28 insertions, 14 deletions

diff --git a/src/IRCServer.py b/src/IRCServer.py
index 65694449..01d03e8d 100644
--- a/src/IRCServer.py
+++ b/src/IRCServer.py
@@ -89,26 +89,17 @@ class Server(IRCObject.Object):
         return self.cached_fileno or self.socket.fileno()
 
     def tls_wrap(self):
-        context = ssl.SSLContext(ssl.PROTOCOL_TLS)
-        context.options |= ssl.OP_NO_SSLv2
-        context.options |= ssl.OP_NO_SSLv3
-        context.options |= ssl.OP_NO_TLSv1
-
-        context.load_default_certs()
-        if self.get_setting("ssl-verify", True):
-            context.verify_mode = ssl.CERT_REQUIRED
-
         client_certificate = self.bot.config.get("tls-certificate", None)
         client_key = self.bot.config.get("tls-key", None)
-        if client_certificate and client_key:
-            context.load_cert_chain(client_certificate, keyfile=client_key)
+        verify = self.get_setting("ssl-verify", True)
 
         server_hostname = None
         if not utils.is_ip(self.connection_params.hostname):
             server_hostname = self.connection_params.hostname
 
-        self.socket = context.wrap_socket(self.socket,
-            server_hostname=server_hostname)
+        self.socket = utils.security.ssl_wrap(self.socket,
+            cert=client_certificate, key=client_key,
+            verify=verify, hostname=server_hostname)
 
     def connect(self):
         ipv4 = self.connection_params.ipv4
diff --git a/src/utils/__init__.py b/src/utils/__init__.py
index c9a42935..188b8e9c 100644
--- a/src/utils/__init__.py
+++ b/src/utils/__init__.py
@@ -1,5 +1,5 @@
 import decimal, io, ipaddress, re, typing
-from src.utils import cli, consts, irc, http, parse
+from src.utils import cli, consts, irc, http, parse, security
 
 TIME_SECOND = 1
 TIME_MINUTE = TIME_SECOND*60
diff --git a/src/utils/security.py b/src/utils/security.py
new file mode 100644
index 00000000..266a767a
--- /dev/null
+++ b/src/utils/security.py
@@ -0,0 +1,23 @@
+import socket, ssl
+
+def ssl_context(cert: str=None, key: str=None, verify: bool=True
+        ) -> ssl.SSLContext:
+    context = ssl.SSLContext(ssl.PROTOCOL_TLS)
+    context.options |= ssl.OP_NO_SSLv2
+    context.options |= ssl.OP_NO_SSLv3
+    context.options |= ssl.OP_NO_TLSv1
+    context.load_default_certs()
+
+    if verify:
+        context.verify_mode = ssl.CERT_REQUIRED
+    if cert and key:
+        context.load_cert_chain(cert, keyfile=key)
+
+    return context
+
+def ssl_wrap(sock: socket.socket, cert: str=None, key: str=None,
+        verify: bool=True, server_side: bool=False, hostname: str=None
+        ) -> ssl.SSLSocket:
+    context = ssl_context(cert=cert, key=key, verify=verify)
+    return context.wrap_socket(sock, server_side=server_side,
+        server_hostname=hostname)