Merge pull request #20 from mszeu/devel

Devel to Masster

Author: mszeu

LICENSE

@@ -658,4 +658,4 @@ specific requirements.
   You should also get your employer (if you work as a programmer) or school,
 if any, to sign a "copyright disclaimer" for the program, if necessary.
 For more information on this, and how to apply and follow the GNU AGPL, see
-<https://www.gnu.org/licenses/>.
+<https://www.gnu.org/licenses/>.

pressureTest.py

@@ -13,7 +13,7 @@
 from typing import Tuple, Dict
 from types import FunctionType
 
-VERSION = "1.1.7"
+VERSION = "1.1.7.2"
 
 
 def decode_n0(response_to_decode: bytes, head_len: int):
@@ -56,7 +56,7 @@ def decode_no(response_to_decode: bytes, head_len: int):
     """
     BUFFER_SIZE: Dict[str, str] = {
         '0': '2K bytes', '1': '8K bytes', '2': '16K bytes', '3': '32K bytes'}
-    NET_PROTO: Dict[str, str] = {'0': 'TCP', '1': 'UDP'}
+    NET_PROTO: Dict[str, str] = {'0': 'UDP', '1': 'TCP'}
     response_to_decode, msg_len, str_pointer = common_parser(response_to_decode, head_len)
     if response_to_decode[str_pointer:str_pointer + 2] == '00':  # No errors
         if len(response_to_decode) >= (24 + head_len):  # Mode 00
@@ -125,8 +125,8 @@ def decode_ni(response_to_decode: bytes, head_len: int):
             str_pointer = str_pointer + 8
             print("Remote port number: ", response_to_decode[str_pointer:str_pointer + 4])
             str_pointer = str_pointer + 4
-            print("Connection Status: ", NET_CONNECTION_STATUS.get(response_to_decode[str_pointer:str_pointer + 1]
-                                                                   , 'Reserved'))
+            print("Connection Status: ", NET_CONNECTION_STATUS.get(response_to_decode[str_pointer:str_pointer + 1],
+                                                                   'Reserved'))
             str_pointer = str_pointer + 1
             print("Duration: ", response_to_decode[str_pointer:str_pointer + 8])
             str_pointer = str_pointer + 8
@@ -729,7 +729,7 @@ def run_test(ip_addr: str, port: int, host_command: str, proto: str = "tcp", hea
     if proto not in ['tcp', 'udp', 'tls']:
         print("invalid protocol parameter, It needs to be tcp, udp or tls")
         return -1
-
+    connection = None
     try:
 
         # calculate the size and format it correctly
@@ -796,18 +796,16 @@ def run_test(ip_addr: str, port: int, host_command: str, proto: str = "tcp", hea
             decoder_funct(data, header_len)
 
     except ConnectionError as e:
-        print("Connection issue: ", e.message)
+        print("Connection issue: ", e)
     except FileNotFoundError as e:
         print("The client certificate file or the client key file cannot be found or accessed.\n" +
-              "Check value passed to the parameters --keyfile and --crtfile", e.message)
+              "Check value passed to the parameters --keyfile and --crtfile", e)
     except Exception as e:
-        if hasattr(e, 'message'):
-            print("Unexpected issue:", e.message)
-        else:
-            print("Unexpected issue:", e)
+        print("Unexpected issue:", e)
 
     finally:
-        connection.close()
+        if connection is not None:
+            connection.close()
 
 
 def common_parser(response_to_decode: bytes, head_len: int) -> Tuple[str, int, int]:
@@ -1016,27 +1014,3 @@ def common_parser(response_to_decode: bytes, head_len: int) -> Tuple[str, int, i
                 run_test(args.host, args.port, command, args.proto, len(args.header), None)
             print("")
         print("DONE")
-
-
-def setup_connection(ip_addr: str, port: int, proto: str = "tcp") -> socket:
-    if proto == "tcp":
-        # creates the TCP socket
-        connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        connection.connect((ip_addr, port))
-        return connection
-
-    elif proto == "tls":
-        # creates the TCP TLS socket
-        connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        ciphers = "ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES128-GCM-SHA256:AES128-GCM-SHA256:AES128-SHA256:HIGH:"
-        ciphers += "!aNULL:!eNULL:!EXPORT:!DSS:!DES:!RC4:!3DES:!MD5:!PSK"
-        ssl_sock = ssl.wrap_socket(connection, args.keyfile, args.crtfile)
-        ssl_sock.connect((ip_addr, port))
-        return connection
-
-    elif proto == "udp":
-        # create the UDP socket
-        connection = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
-        return connection
-
-    return None