1 files changed, 26 insertions, 12 deletions

diff --git a/server/mfad.py b/server/mfad.py
index 17d2585..18a048a 100755
--- a/server/mfad.py
+++ b/server/mfad.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 import socket
+import ssl
 import os
 import sys
 import time
@@ -211,8 +212,9 @@ def parse_pam_data(data):
 def handle_pam(db, conn, addr):
     # Get request and data from PAM module
     header = conn.recv(HEADER_LENGTH).decode(FORMAT)
-    if header == "":
-        die("error: lost connection to pam module")
+    if len(header) != HEADER_LENGTH:
+        conn.close()
+        die("error: invalid data from PAM module")
     data_length = int(header)
     pam_data = conn.recv(data_length).decode(FORMAT)
     print("Got pam_data: " + pam_data)
@@ -238,19 +240,31 @@ def handle_pam(db, conn, addr):
 
 
 def listen_client(db, addr, port):
-    with socket.create_server((addr, port)) as server:
-        while True:
-            conn, addr = server.accept()
-            thread = threading.Thread(target=handle_client,args=(db, conn,addr))
-            thread.start()
+    context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+    context.load_cert_chain(certfile="server/cert.pem", keyfile="server/key.pem")
+    with socket.create_server((addr, port)) as sock:
+        with context.wrap_socket(sock, server_side=True) as tls_socket:
+            while True:
+                try:
+                    conn, addr = tls_socket.accept()
+                    thread = threading.Thread(target=handle_client,args=(db,conn,addr))
+                    thread.start()
+                except ssl.SSLError:
+                    print("client: ssl handshake error")
 
 
 def listen_pam(db, addr, port):
-    with socket.create_server((addr,port)) as pam_server:
-        while True:
-            conn, addr = pam_server.accept()
-            thread = threading.Thread(target=handle_pam,args=(db, conn,addr))
-            thread.start()
+    context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
+    context.load_cert_chain(certfile="server/cert.pem", keyfile="server/key.pem")
+    with socket.create_server((addr,port)) as sock:
+        with context.wrap_socket(sock, server_side=True) as tls_socket:
+            while True:
+                try:
+                    conn, addr = tls_socket.accept()
+                    thread = threading.Thread(target=handle_pam,args=(db, conn,addr))
+                    thread.start()
+                except ssl.SSLError:
+                    print("pam: ssl handshake error")
 
 
 ################################################################################