Преглед изворни кода

Refactored field access in Sample.

Sample now uses a getter to acess properties.

Properties for high, mid, and low activities use a prefix seperated by '.'
Thomas Flucke пре 6 година
родитељ
комит
a7a98e2f34
1 измењених фајлова са 41 додато и 5 уклоњено
  1. 41 5
      src/feature-extractor/sample.py

+ 41 - 5
src/feature-extractor/sample.py

@@ -37,6 +37,7 @@ class Sample:
             .total_seconds()
         
     def __init__(self, packets, keylog: typing.TextIO):
+        self.__general = {}
         self.__extract_tag(keylog)
         self.__extract_activity_stats(packets)
         self.__extract_packet_stats(packets)
@@ -44,9 +45,9 @@ class Sample:
     def __extract_tag(self, keylog: typing.TextIO):
         import os
         dir_guided = os.path.dirname(keylog)
-        self.is_guided = os.path.basename(dir_guided) == "y"
+        self["is_guided"] = os.path.basename(dir_guided) == "y"
         dir_user = os.path.dirname(dir_guided)
-        self.user = os.path.basename(dir_user)
+        self["user"] = os.path.basename(dir_user)
 
     def __extract_activity_stats(self, packets):
         high_activity = []
@@ -65,7 +66,7 @@ class Sample:
                 high_activity.append(q)
             while Sample.__packet_time(q[0]) + self.lookback < ptime:
                 q = q[1:]
-        self.activities = {
+        self.__activities = {
             "high": Sample.__count_activity_stats(high_activity),
             "mid": Sample.__count_activity_stats(mid_activity),
             "low": Sample.__count_activity_stats(low_activity)
@@ -73,10 +74,45 @@ class Sample:
 
     def __count_activity_stats(arr):
         return {
-            "total packets": len(arr)
+            "total_packets": len(arr)
         }
         
     def __extract_packet_stats(self, pcap):
         start = Sample.__packet_time(pcap[0])
         end = Sample.__packet_time(pcap[-1])
-        self.average_iat = (end - start) / len(pcap)
+        self["average_iat"] = (end - start) / len(pcap)
+
+    def __is_valid_prefix(pre):
+        return pre in ["high", "mid", "low"]
+        
+    def __setitem__(self, key, value):
+        vals = key.split(".")
+        if 1 == len(vals):
+            self.__general[key] = value
+        elif 2 == len(vals) and Sample.__is_valid_prefix(vals[0]):
+            self.__activities[vals[0]][vals[1]] = value
+        else:
+            raise ValueError('Unable to access value at %s, unknown prefix.' % key)
+
+    def __contains__(self, key):
+        vals = key.split(".")
+        if 1 == len(vals):
+            return True
+        elif 2 == len(vals):
+            return Sample.__is_valid_prefix(vals[0])
+        else:
+            return False
+
+    def __getitem__(self, key):
+        vals = key.split(".")
+        if 1 == len(vals):
+            return self.__general[key]
+        elif 2 == len(vals) and Sample.__is_valid_prefix(vals[0]):
+            return self.__activities[vals[0]][vals[1]]
+        else:
+            raise ValueError('Unable to access value at %s, unknown prefix.' % key)
+    
+    def __str__(self):
+        return "Sample: {%s, high: %s, mid: %s, low: %s}" % \
+            (self.__general, self.__activities["high"],
+             self.__activities["mid"], self.__activities["low"])