Thrift chang to default ssl context

Instead of the filepath of the certs the default ssl context is used
Instead of the three file paths you can give as 6th arguments true / false to enable / disable ssl

1 files changed, 11 insertions, 8 deletions

diff --git a/main.py b/main.py
index 4811db2..3008c39 100755
--- a/main.py
+++ b/main.py
@@ -3,6 +3,7 @@ import datetime
 import logging
 import os
 import shutil
+import ssl
 import subprocess
 import tftpy
 import thriftpy
@@ -117,20 +118,26 @@ def tftp(name, host, port, filename):
     logStatus(statusList[-1])
   
 # Creates the thrift client and prwlp-pxe.ruf.uni-freiburg.de .ceeds the get Organisations call. Server can either be SAT or MASTER 
-def thrift(name, ip, port, server, cafile=None, certfile=None, keyfile=None):
+def thrift(name, ip, port, server, SSL=False):
   host = ip + ':' + str(port)
   print('THRIFT request ' + host + ' ...', end='')
   bwlp_thrift = thriftpy.load('bwlp.thrift', module_name='bwlp_thrift')
+  
+  # SSL factory
+  ssl_factory = None
+  if SSL == 'true':
+    # ssl_factory = ssl.create_default_context()
+    ssl_factory = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
 
   organisations = []
   # Different clients for SAT / Master is needed.
   try:
     if server == 'SAT':
-      satserver = make_client(bwlp_thrift.SatelliteServer, host=ip, port=port, trans_factory=TFramedTransportFactory(), cafile=cafile, certfile=certfile, keyfile=keyfile)
+      satserver = make_client(bwlp_thrift.SatelliteServer, host=ip, port=port, trans_factory=TFramedTransportFactory(), ssl_context=ssl_factory)
       organisations = satserver.getAllOrganizations()
 
     elif server == 'MASTER':
-      masterserver = make_client(bwlp_thrift.MasterServer, host=ip, port=port, trans_factory=TFramedTransportFactory(), cafile=cafile, certfile=certfile, keyfile=keyfile)
+      masterserver = make_client(bwlp_thrift.MasterServer, host=ip, port=port, trans_factory=TFramedTransportFactory(), ssl_context=ssl_factory)
       organisations = masterserver.getOrganizations()
 
     organisationList = []
@@ -211,11 +218,7 @@ def check(entry):
     https(entry[1], entry[2])
   elif entry[0] == 'thrift':
     if (len(entry) < 5): return
-    elif (len(entry) == 5):
-      thrift(entry[1], entry[2], int(entry[3]), entry[4])
-    elif (len(entry) == 8):
-      # SSL attributes
-      thrift(entry[1], entry[2], int(entry[3]), entry[4], entry[5], entry[6], entry[7])
+    thrift(entry[1], entry[2], int(entry[3]), entry[4], entry[5] if len(entry) == 6 else 'false')
   elif entry[0] == 'tftp':
     if (len(entry) < 5): return
     tftp(entry[1], entry[2], int(entry[3]), entry[4])