Add thrift ssl support, more error handling

Thrift ssl support:
Extend your thrift bwlp.config with the path of cafile, certfile and keyfile (see example bwlp.config)
Add ssl file not found error handling
Add error handling for all exceptions and log them

2 files changed, 20 insertions, 8 deletions

diff --git a/bwlp.config b/bwlp.config
index 207f2ca..9a9ce84 100644
--- a/bwlp.config
+++ b/bwlp.config
@@ -2,6 +2,7 @@
 # https	name	host
 # thrift	name	host	port	SAT/MASTER
 # tftp	name	host	port	filename
+# tftp	name	host	port	filename	cafile	certfile	keyfile
 https	Masterserver HTTPS Service	https://bwlp-masterserver.ruf.uni-freiburg.de
 https	Fail Test	5121236
 https	Fail Test	https://www.amazony.de/
@@ -11,10 +12,11 @@ ping	Backup fileserver Ping	bwlp-backup.ruf.uni-freiburg.de
 ping	Ping Fail	0.0.0.1
 thrift	Thrift SAT	132.230.8.192	9090	SAT
 thrift	Thrift MASTER	132.230.4.16	9090	MASTER
-thrift	Thrift MASTER	132.230.4.16	9091	MASTER
+thrift	Thrift Timeout	132.230.4.16	9092	MASTER
+thrift	Thrift SSL test	132.230.4.16	9091	MASTER	ssl/CA.pem	ssl/client.crt	ssl/client.key
 ping	BAS Ping	bas.intra.uni-freiburg.de
 ping	yc	127.0.0.2
 ping	yx	127.0.0.3
 tftp	TFTP bwlp lpxelinux	bwlp-pxe.ruf.uni-freiburg.de	69	lpxelinux.0
 tftp	TFTP bas ipxe	bas.intra.uni-freiburg.de	69	ipxe.0
-tftp	TFTP bas fail test	bas.intra.uni-freiburg.de	69	failtest.0
+tftp	TFTP bas fail test	bas.intra.uni-freiburg.de	69	failtest.0
\ No newline at end of file
diff --git a/main.py b/main.py
index ea13990..48c21da 100755
--- a/main.py
+++ b/main.py
@@ -116,7 +116,7 @@ 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):
+def thrift(name, ip, port, server, cafile=None, certfile=None, keyfile=None):
   host = ip + ':' + str(port)
   print('THRIFT request ' + host + ' ...', end='')
   bwlp_thrift = thriftpy.load('bwlp.thrift', module_name='bwlp_thrift')
@@ -125,11 +125,11 @@ def thrift(name, ip, port, server):
   # Different clients for SAT / Master is needed.
   try:
     if server == 'SAT':
-      satserver = make_client(bwlp_thrift.SatelliteServer, ip, port, trans_factory=TFramedTransportFactory())
+      satserver = make_client(bwlp_thrift.SatelliteServer, host=ip, port=port, trans_factory=TFramedTransportFactory(), cafile=cafile, certfile=certfile, keyfile=keyfile)
       organisations = satserver.getAllOrganizations()
 
     elif server == 'MASTER':
-      masterserver = make_client(bwlp_thrift.MasterServer, ip, port, trans_factory=TFramedTransportFactory())
+      masterserver = make_client(bwlp_thrift.MasterServer, host=ip, port=port, trans_factory=TFramedTransportFactory(), cafile=cafile, certfile=certfile, keyfile=keyfile)
       organisations = masterserver.getOrganizations()
 
     organisationList = []
@@ -142,7 +142,13 @@ def thrift(name, ip, port, server):
       statusList.append(Status(name, host, 'success', 'Online (' + str(len(organisationList)) + ')', 'thrift', msg=str(len(organisationList)) + " organizations found",data=organisationList))
     print('\033[92m' + 'success' + '\033[0m')
   except ConnectionResetError:
-    statusList.append(Status(name, host, 'error', 'Offline', 'thrift', msg="ConnectionResetError: [Errno 104] Connection reset by peer"))
+    statusList.append(Status(name, host, 'error', 'Offline', 'thrift', msg='ConnectionResetError: [Errno 104] Connection reset by peer'))
+    print('\033[91m' + 'error' + '\033[0m')
+  except FileNotFoundError:
+    statusList.append(Status(name, host, 'error', 'Offline', 'thrift', msg='SSL path incorrect (FileNotFound)'))
+    print('\033[91m' + 'error' + '\033[0m')
+  except Exception as e:
+    statusList.append(Status(name, host, 'error', 'Offline', 'thrift', msg=str(e)))
     print('\033[91m' + 'error' + '\033[0m')
   finally:
     logStatus(statusList[-1])
@@ -203,7 +209,11 @@ def check(entry):
     https(entry[1], entry[2])
   elif entry[0] == 'thrift':
     if (len(entry) < 5): return
-    thrift(entry[1], entry[2], int(entry[3]), entry[4])
+    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])
   elif entry[0] == 'tftp':
     if (len(entry) < 5): return
     tftp(entry[1], entry[2], int(entry[3]), entry[4])
@@ -261,7 +271,7 @@ for entry in reversed(logEntries[cut_amount:]):
           + entry.msg + '</div></div></div>')
 
 # Replace the %CONTENT% in the template with the actual html code.
-now = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
 html = open('template.html')
 html2 = open('status.html', 'w')
 for line in html: