Fix/297 (#302)

Fixed some Codacy-related bugs. Additionally, this should fix issue #297

DeviceClassifier/OneLayer/train_OneLayer.py

@@ -5,7 +5,6 @@
 ('models/OneLayerModel' by default).
 '''
 import argparse
-import sys
 
 from poseidonml.config import get_config
 from poseidonml.Model import Model

DeviceClassifier/RandomForest/train_RandomForest.py

@@ -5,7 +5,6 @@
 ('models/RandomForestModel' by default).
 '''
 import argparse
-import sys
 
 from poseidonml.config import get_config
 from poseidonml.Model import Model

utils/common.py

@@ -273,7 +273,7 @@ def update_data(
             update_list = ast.literal_eval(
                 updates[b'timestamps'].decode('ascii'))
             self.logger.debug('Got previous updates from %s', source_mac)
-        except Exception as e:
+        except:
             self.logger.debug('No previous updates found for %s', source_mac)
             update_list = []
 
@@ -288,8 +288,8 @@ def update_data(
         try:
             self.r.hmset(source_mac, redis_times)
             self.r.sadd('mac_addresses', source_mac)
-        except Exception as e:
-            self.logger.debug('Could not store update time')
+        except (ConnectionError, TimeoutError) as e:
+            self.logger.debug('Could not store update time because: %s', str(e))
 
         return key
 
@@ -351,5 +351,5 @@ def get_config(self):
             #self.batch_size = config['batch size']
         except Exception as e:  # pragma: no cover
             self.logger.error(
-                "unable to read 'opts/config.json' properly because: %s", str(e))
+                "Unable to read 'opts/config.json' properly because: %s", str(e))
         return

utils/featurizer.py

@@ -33,13 +33,9 @@ def extract_features(session_dict, capture_source=None, max_port=None):
     '''
 
     # Get featurization info from config
-    try:
-        config = get_config()
-        address_type = config['source identifier']
-        if max_port is None:
-            max_port = config['max port']
-    except Exception as e:
-        address_type = 'MAC'
+    config = get_config()
+    address_type = config.get('source identifier', 'MAC')
+    max_port = config.get('max port', max_port) if not max_port else None
 
     # If the capture source isn't specified, default to the most used address
     if capture_source is None:

utils/pcap_utils.py

@@ -193,9 +193,8 @@ def packet_size(packet):
     size = packet[1][32:36]
     try:
         size = int(size, 16)
-    except Exception as e:
+    except ValueError:
         size = 0
-
     return size
 
 
@@ -300,8 +299,11 @@ def clean_session_dict(sessions, source_address=None):
     def clean_dict(sessions, source_address):
         cleaned_sessions = OrderedDict()
         for key, packets in sessions.items():
-            address_1, port_1 = get_ip_port(key[0])
-            address_2, port_2 = get_ip_port(key[1])
+            # TODO: Removing port_1 and port_2 (i.e., returned val [1])
+            # due to unuse, but I'm a little surprised we aren't using
+            # this... O_o
+            address_1 = get_ip_port(key[0])[0]
+            address_2 = get_ip_port(key[1])[0]
 
             first_packet = sessions[key][0][1]
             source_mac, destination_mac = extract_macs(first_packet)

utils/reader.py

@@ -115,6 +115,7 @@ def packetizer(path):
 
     # Read get the pcap info with tcpdump
     FNULL = open(os.devnull, 'w')
+    # TODO: yikes @ the shell=True + unvalidated user input
     proc = subprocess.Popen(
         'tcpdump -nn -tttt -xx -r' + path,
         shell=True,
@@ -166,11 +167,10 @@ def sessionizer(path, duration=None, threshold_time=None):
 
     # Get threshold time from config
     if threshold_time is None:
-        try:
-            config = get_config()
-            threshold_time = config['session threshold']
-        except Exception as e:
-            threshold_time = 120
+        # TODO: error-check threshold_time, if < 0
+        # error-check all config parameters, as well.
+        config = get_config()
+        threshold_time = config.get('session threshold', 120)
 
     for head, packet in packet_dict.items():
         time = head[0]