Make Thirft and Ping handlers support multiple DNS recordsHEADmaster

1 files changed, 103 insertions, 52 deletions

diff --git a/main.py b/main.py
index 5e38c3e..d0b47df 100755
--- a/main.py
+++ b/main.py
@@ -1,8 +1,10 @@
 #!/usr/bin/python3
+import traceback
 import datetime
 import logging
 import os
 import shutil
+import socket
 import ssl
 import subprocess
 import time
@@ -59,27 +61,38 @@ class Organisation:
 def ping(name, hostname):
   # Ping a hostname and tell if the server is up or down.
   print('Ping request ' + hostname + ' ... ', end='')
-  response = subprocess.Popen(
-    ['ping', '-c', '1', hostname],
-    stdout=subprocess.PIPE,
-    stderr=subprocess.PIPE
-  )
-  response.wait()
-  exit_code = response.returncode
-  response = response.communicate()
-  error = response[1].decode('utf-8')
-  # Error happend
-  if error != '' or exit_code != 0:
-    err_msg = ''
-    if exit_code == 1:
-      err_msg = 'Destination Host Unreachable'
-    elif exit_code == 2:
-      err_msg = error.rstrip()
-    print('\033[91m' + 'error' + '\033[0m')
-    statusList.append(Status(name, hostname, 'error', 'Offline', 'ping', err_msg))
-  else:
-    print('\033[92m' + 'success' + '\033[0m')
-    statusList.append(Status(name, hostname, 'success', 'Online', 'ping'))
+  try:
+    ips = socket.gethostbyname_ex(hostname)[2]
+    if len(ips) == 0:
+      statusList.append(Status(name, hostname, 'error', 'Offline', 'thrift', msg='No IPs for ' + hostname))
+  except socket.gaierror as e:
+    print(f"DNS resolution failed: {e}")
+    statusList.append(Status(name, hostname, 'error', 'Offline', 'thrift', msg='DNS for ' + hostname + ' failed: ' + str(e)))
+    ips = []
+
+  for ip in ips:
+    response = subprocess.Popen(
+      ['ping', '-w', '2', '-c', '1', ip],
+      stdout=subprocess.PIPE,
+      stderr=subprocess.PIPE
+    )
+    response.wait()
+    exit_code = response.returncode
+    response = response.communicate()
+    error = response[1].decode('utf-8')
+    # Error happend?
+    if error != '' or exit_code != 0:
+      err_msg = ''
+      if exit_code == 1:
+        err_msg = 'Destination Host Unreachable'
+      elif exit_code == 2:
+        err_msg = error.rstrip()
+      print('\033[91m' + 'error' + '\033[0m')
+      statusList.append(Status(name, hostname, 'error', 'Offline', 'ping', err_msg))
+    else: # Success
+      print('\033[92m' + 'success' + '\033[0m')
+      statusList.append(Status(name, hostname, 'success', 'Online', 'ping'))
+      break
     
   logStatus(statusList[-1])
 
@@ -121,19 +134,41 @@ 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, SSL=False):
-  host = ip + ':' + str(port)
+def thrift(name, hostname, port, server, SSL=False):
+  host = hostname + ':' + str(port)
   print('THRIFT request ' + host + ' ...', end='')
   try:
     import thriftpy
-    from thriftpy.rpc import make_client
+    from thriftpy.transport import TSocket
+    from thriftpy.transport import TTransportBase
     from thriftpy.transport import TFramedTransportFactory
+    from thriftpy.protocol import TBinaryProtocolFactory
+    from thriftpy.thrift import TClient
+  except Exception as e:
+    print('\033[91m' + 'error' + '\033[0m')
+    statusList.append(Status(name, host, 'error', 'Offline', 'thrift', msg=str(e)))
+    logStatus(statusList[-1])
+    return
+
+  try:
     bwlp_thrift = thriftpy.load('bwlp.thrift', module_name='bwlp_thrift')
   except Exception as e:
     print('\033[91m' + 'error' + '\033[0m')
+    traceback.print_exc()
     statusList.append(Status(name, host, 'error', 'Offline', 'thrift', msg=str(e)))
     logStatus(statusList[-1])
     return
+
+  if server == 'SAT':
+    service = bwlp_thrift.SatelliteServer
+  elif server == 'MASTER':
+    service = bwlp_thrift.MasterServer
+  else:
+    print('\033[91m' + 'error' + '\033[0m')
+    traceback.print_exc()
+    statusList.append(Status(name, host, 'error', 'Offline', 'thrift', msg='Unknown service type: ' + server))
+    logStatus(statusList[-1])
+    return
   
   # SSL factory
   ssl_factory = None
@@ -142,37 +177,53 @@ def thrift(name, ip, port, server, SSL=False):
     ssl_factory = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
     ssl_factory.load_default_certs()
 
-  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(), ssl_context=ssl_factory)
-      organisations = satserver.getAllOrganizations()
-
-    elif server == 'MASTER':
-      masterserver = make_client(bwlp_thrift.MasterServer, host=ip, port=port, trans_factory=TFramedTransportFactory(), ssl_context=ssl_factory)
-      organisations = masterserver.getOrganizations()
+    ips = socket.gethostbyname_ex(hostname)[2]
+    if len(ips) == 0:
+      statusList.append(Status(name, host, 'error', 'Offline', 'thrift', msg='No IPs for ' + hostname))
+  except socket.gaierror as e:
+    print(f"DNS resolution failed: {e}")
+    statusList.append(Status(name, host, 'error', 'Offline', 'thrift', msg='DNS lookup failed for ' + hostname + ': ' + str(e)))
+    ips = []
+
+  # Try them all
+  organisationList = []
+  for ip in ips:
+    host = hostname + ':' + str(port)
+    organisations = []
+    try:
+      #sock = CustomSSLTSocket(ip, port, hostname, ssl_factory)
+      sock = socket.create_connection((ip, port))
+      if ssl_factory is not None:
+        sock = ssl_factory.wrap_socket(sock, server_hostname=hostname)
+      client = TClient(service, TBinaryProtocolFactory().get_protocol(TFramedTransportFactory().get_transport(TSocket(sock=sock))))
+
+      if server == 'SAT':
+        organisations = client.getAllOrganizations()
+      elif server == 'MASTER':
+        organisations = client.getOrganizations()
+
+      organisationList = []
+      for org in organisations:
+        organisationList.append(Organisation(org.organizationId, org.displayName))
+
+      if len(organisationList) != 0:
+        print('\033[92m' + 'success' + '\033[0m')
+        statusList.append(Status(name, host, 'success', 'Online (' + str(len(organisationList)) + ')', 'thrift', msg=str(len(organisationList)) + " organizations found",data=organisationList))
+        break
+    except ConnectionResetError:
+      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')
+      traceback.print_exc()
 
-    organisationList = []
-    for org in organisations:
-      organisationList.append(Organisation(org.organizationId, org.displayName))
+  logStatus(statusList[-1])
 
-    if len(organisationList) == 0:
-      statusList.append(Status(name, host, 'warning', 'Online (' + str(len(organisationList)) + ')', 'thrift', msg=str(len(organisationList)) + " organizations found",data=organisationList))
-    else:
-      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'))
-    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])
 
 # Parses the log from the logfile. Fills the logEntries array.
 def parseLog():