main in funktionen aufgeteilt

setstatus.py

@@ -17,6 +17,7 @@ import ssl
 import socket
 import logging
 import configparser
+from time import sleep
 from sys import exit, argv
 
 
@@ -53,7 +54,7 @@ def check_arguments(argv):
 
 def log_config(config):
     '''
-    Logs the config with level debug.
+    Logs the config if loglevel is debug.
     '''
     logging.debug('Using config:')
     for section in config.sections():
@@ -61,6 +62,76 @@ def log_config(config):
         for i in config[section]:
             logging.debug('  {}: {}'.format(i, config[section][i]))
 
+def create_ssl_context(config):
+    '''
+    '''
+    context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH,
+                                         cafile=config['server']['cert'])
+    if not context:
+        logging.error('Failed to create SSL Context')
+        return False
+    context.set_ciphers('EECDH+AESGCM') # only ciphers for tls 1.2 and 1.3
+    context.options |= getattr(ssl._ssl, 'OP_NO_COMPRESSION', 0)
+    try:
+        context.load_cert_chain(certfile=config['client']['cert'],
+                                keyfile=config['client']['key'])
+    except Exception as e:
+        logging.error('Failed to load cert chain')
+        return False;
+    logging.debug('SSL context created')
+    return context
+
+def create_ssl_socket(config, context):
+    '''
+    '''
+    bare_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
+    if not bare_socket:
+        logging.error('Socket creation failed')
+        return False
+    logging.debug('Socket created')
+    try:
+        secure_socket = context.wrap_socket(bare_socket, server_side=False,
+                                   server_hostname=config['server']['fqdn'])
+        logging.debug('Socket wrapped with SSL')
+    except Exception as e:
+        logging.error('Context wrapper failed: {}'.format(e))
+        return False
+    try:
+        secure_socket.settimeout(float(config['general']['timeout']))
+    except Exception as e:
+        logging.debug('Failed to set timeout: {}'.format(e))
+    return secure_socket
+
+def create_ssl_connection(config, context):
+    '''
+    '''
+    counter = 0
+    conn = False
+
+    while conn is False and counter < 5:
+        ssl_socket = create_ssl_socket(config, context)
+        if ssl_socket is False:
+            exit(4)
+        try:
+            logging.debug('Connect {}: {}'.format(config['server']['host'],
+                                                  int(config['server']['port'])))
+            conn = ssl_socket.connect((config['server']['host'],
+                                      int(config['server']['port'])))
+        except socket.timeout:
+            logging.error('Connection timeout')
+            ssl_socket.close()
+            sleep(5)
+            counter += 1
+        except Exception as e:
+            logging.error('Connection failed: {}'.format(e))
+            exit(5)
+    logging.debug('Conection established')
+    logging.debug('Peer certificate commonName: {}'.format(
+                  ssl_socket.getpeercert()['subject'][5][0][1]))
+    logging.debug('Peer certificate serialNumber: {}'.format(
+                  ssl_socket.getpeercert()['serialNumber']))
+    return ssl_socket
+
 
 def main():
 
@@ -115,53 +186,29 @@ def main():
              config['client']['cert'],
              config['client']['key'])
     if check_certs(certs) is False:
-        sys.exit(1)
+        exit(1)
 
     # check cli arguments
     STATUS = check_arguments(argv)
     if STATUS is None:
         logging.error('No valid status given')
-        sys.exit(2)
+        exit(2)
 
     # create ssl context
-    context = ssl.create_default_context(ssl.Purpose.SERVER_AUTH,
-                                         cafile=config['server']['cert'])
-    context.set_ciphers('EECDH+AESGCM') # only ciphers for tls 1.2 and 1.3
-    context.options |= getattr(ssl._ssl, 'OP_NO_COMPRESSION', 0)
-    context.load_cert_chain(certfile=config['client']['cert'],
-                            keyfile=config['client']['key'])
-    logging.debug('SSL context created')
-
-    # create socket and send status bit
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) as mySocket:
-        logging.debug('Socket created')
-        try:
-            conn = context.wrap_socket(mySocket, server_side=False,
-                                       server_hostname=config['server']['fqdn'])
-            logging.debug('Connection wrapped with ssl.context')
-        except Exception as e:
-            logging.error('Context wrapper failed: {}'.format(e))
-        try:
-            conn.settimeout(float(config['general']['timeout']))
-        except Exception as e:
-            logging.debug('Failed to set timeout: {}'.format(e))
-        try:
-            conn.connect((config['server']['host'], int(config['server']['port'])))
-        except socket.timeout:
-            logging.error('Connection timeout')
-        except Exception as e:
-            logging.error('Connection failed: {}'.format(e))
+    context = create_ssl_context(config)
+    if context is False:
         exit(3)
-        logging.debug('Peer certificate commonName: {}'.format(
-                      conn.getpeercert()['subject'][5][0][1]))
-        logging.debug('Peer certificate serialNumber: {}'.format(
-                      conn.getpeercert()['serialNumber']))
+
+    # get connection
+    conn = create_ssl_connection(config, context)
+
+    # send status
     try:
         logging.debug('Send new status: {}'.format(STATUS))
         conn.send(STATUS)
     except Exception as e:
         logging.error('Error: {}'.format(e))
-            exit(4)
+        exit(6)
     try:
         RESPONSE = conn.recv(1)
         logging.debug('Server returns: {}'.format(RESPONSE))
@@ -172,7 +219,7 @@ def main():
         logging.debug('Disconnect from server')
     except Exception as e:
         logging.error('Error: {}'.format(e))
-            exit(5)
+        exit(7)
 
 
 if __name__ == '__main__':