'Refactored by Sourcery'

analysis/classifier.py

@@ -64,12 +64,9 @@ def test(self, data, labels=None):
 
         out = list(self.logreg.predict_proba(data))
 
-        if labels == None:
+        if labels is None:
             return out, None, None
-        correct = 0
-        for i in range(len(out)):
-            if self.get_label(out[i]) == labels[i]:
-                correct += 1
+        correct = sum(self.get_label(out[i]) == labels[i] for i in range(len(out)))
         acc = correct/float(len(out))
 
         return out, acc, correct
@@ -81,11 +78,7 @@ def get_parameters(self):
         return self.logreg.coef_, self.logreg.intercept_
 
     def get_num_nonzero_params(self):
-        tmp = 0
-        for x in self.logreg.coef_[0]:
-            if x != 0.0:
-                tmp += 1
-        return tmp
+        return sum(x != 0.0 for x in self.logreg.coef_[0])
 
     def CV(self, data, labels, folds=10):
         self.is_cv = True
@@ -104,15 +97,15 @@ def CV(self, data, labels, folds=10):
         for i in range(folds):
             start = int((i/float(folds))*len(data))
             end = int(((i+1)/float(folds))*len(data))
-            train_data = data[0:start]+data[end:]
-            train_labels = labels[0:start]+labels[end:]
+            train_data = data[:start] + data[end:]
+            train_labels = labels[:start] + labels[end:]
             test_data = data[start:end]
             test_labels = labels[start:end]
             self.train(train_data,train_labels)
 #            self.train(train_data+test_data,train_labels+test_labels)
             out,acc,correct = self.test(test_data,test_labels)
             self.non_zero_params.append(self.get_num_nonzero_params())
-            
+
             out_cv.extend(out)
             acc_cv.append(acc)
             correct_cv.append(correct)

analysis/data_parser.py

@@ -156,28 +156,24 @@ def getIndividualFlowIPTs(self):
             return None
 
         data = []
-        if self.compact:
-            numRows = 10
-            binSize = 50.0
-        else:
-            numRows = 30
-            binSize = 50.0
+        numRows = 10 if self.compact else 30
+        binSize = 50.0
         for flow in self.flows:
             transMat = np.zeros((numRows,numRows))
             if len(flow['packets']) == 0:
                 continue
             elif len(flow['packets']) == 1:
-                curIPT = min(int(flow['packets'][0]['ipt']/float(binSize)),numRows-1)
+                curIPT = min(int(flow['packets'][0]['ipt'] / binSize), numRows-1)
                 transMat[curIPT,curIPT] = 1
                 data.append(list(transMat.flatten()))
                 continue
 
             # get raw transition counts
             for i in range(1,len(flow['packets'])):
-                prevIPT = min(int(flow['packets'][i-1]['ipt']/float(binSize)),numRows-1)
-                curIPT = min(int(flow['packets'][i]['ipt']/float(binSize)),numRows-1)
+                prevIPT = min(int(flow['packets'][i-1]['ipt'] / binSize), numRows-1)
+                curIPT = min(int(flow['packets'][i]['ipt'] / binSize), numRows-1)
                 transMat[prevIPT,curIPT] += 1
-                
+
             # get empirical transition probabilities
             for i in range(numRows):
                 if float(np.sum(transMat[i:i+1])) != 0:
@@ -200,14 +196,14 @@ def getIndividualFlowMetadata(self):
 
             key = flow['sa'].replace('.','')+flow['da'].replace('.','')+str(flow['sp'])+str(flow['dp'])+str(flow['pr'])
 
-            if flow['dp'] != None:
-                tmp.append(float(flow['dp'])) # destination port
-            else:
+            if flow['dp'] is None:
                 tmp.append(0) # ICMP/etc.
-            if flow['sp'] != None:
-                tmp.append(float(flow['sp'])) # source port
             else:
+                tmp.append(float(flow['dp'])) # destination port
+            if flow['sp'] is None:
                 tmp.append(0) # ICMP/etc.
+            else:
+                tmp.append(float(flow['sp'])) # source port
             if 'num_pkts_in' in flow:
                 tmp.append(flow['num_pkts_in']) # inbound packets
             else:
@@ -228,13 +224,9 @@ def getIndividualFlowMetadata(self):
             if flow['packets'] == []:
                 tmp.append(0)
             else:
-                time = 0
-                for packet in flow['packets']:
-                    time += packet['ipt']
+                time = sum(packet['ipt'] for packet in flow['packets'])
                 tmp.append(time)
 
             data.append(tmp)
 
-        if data == []:
-            return None
-        return data
+        return data or None

fingerprinting/fingerprint_ui.py

@@ -74,9 +74,12 @@ def fingerprint(pcap_file=None):
     if pcap_file == 'upload' and request.files.get('upload') != None:
         upload = request.files.get('upload')
         dir_name = tempfile.mkdtemp()
-        upload.save(dir_name + 'temp.pcap')
+        upload.save(f'{dir_name}temp.pcap')
+
+        fps_ = fingerprinter.extract_fingerprints(
+            [f'{dir_name}temp.pcap'], detailed=True
+        )
 
-        fps_ = fingerprinter.extract_fingerprints([dir_name + 'temp.pcap'], detailed=True)
 
     # clean up temporary directories
     if dir_name != None and os.path.isdir(dir_name):
@@ -102,6 +105,6 @@ def client_info(client_ip):
 
 @route('/static/<filename:path>')
 def send_static(filename):
-    return static_file(filename, root=dir_path+'/static/')
+    return static_file(filename, root=f'{dir_path}/static/')
 
 run(host='localhost', port=8080)

fingerprinting/fingerprinter.py

@@ -61,7 +61,7 @@ def __init__(self, database, port=None, output=None):
         self.port = port
         self.tls = TLSFingerprint(database)
 
-        if output == None:
+        if output is None:
             self.out_file_pointer = None
         elif output == sys.stdout:
             self.out_file_pointer = sys.stdout
@@ -103,10 +103,10 @@ def lookup_fingerprint_string(self, fp_str):
         if fp_str in self.tls.fp_db:
             fp_ = copy.deepcopy(self.tls.fp_db[fp_str])
             del fp_['tls_features']['cs_mapping']
-            self.write_record(fp_)
         else:
             fp_ = self.tls.gen_unknown_fingerprint(fp_str, False)
-            self.write_record(fp_)
+
+        self.write_record(fp_)
 
 
     def extract_fingerprints(self, input_files, detailed=False):
@@ -138,20 +138,19 @@ def extract_fingerprints(self, input_files, detailed=False):
                         break # no data error?
                     ip = eth.data
 
-                    if (type(ip) != dpkt.ip.IP and type(ip) != dpkt.ip6.IP6) or type(ip.data) != dpkt.tcp.TCP:
+                    if (
+                        type(ip) not in [dpkt.ip.IP, dpkt.ip6.IP6]
+                        or type(ip.data) != dpkt.tcp.TCP
+                    ):
                         continue
 
                     tcp = ip.data
                     data = tcp.data
 
-                    if self.port != None:
-                        if not tcp.dport == int(self.port):
-                            continue
+                    if self.port != None and tcp.dport != int(self.port):
+                        continue
 
-                    if type(ip) == dpkt.ip.IP:
-                        add_fam = socket.AF_INET
-                    else:
-                        add_fam = socket.AF_INET6
+                    add_fam = socket.AF_INET if type(ip) == dpkt.ip.IP else socket.AF_INET6
                     flow_key = (socket.inet_ntop(add_fam,ip.src), tcp.sport, socket.inet_ntop(add_fam,ip.dst), tcp.dport)
 
 
@@ -175,7 +174,7 @@ def extract_fingerprints(self, input_files, detailed=False):
                 if capture_type == 'offline':
                     break
 
-        if self.out_file_pointer != None and self.out_file_pointer != sys.stdout:
+        if self.out_file_pointer not in [None, sys.stdout]:
             self.out_file_pointer.close()
 
         return data_