#66 Attack recognition improvement

Merged
carlos.garcia merged 3 commits from stefan.schmidt/attack_recognition into SPIN/master 6 years ago
2 changed files with 65 additions and 0 deletions
  1. 51 0
      code/Core/AttackController.py
  2. 14 0
      code/ID2TLib/Utility.py

+ 51 - 0
code/Core/AttackController.py

@@ -1,11 +1,13 @@
 import importlib
 import importlib
 import sys
 import sys
+import difflib
 
 
 import Attack.AttackParameters as atkParam
 import Attack.AttackParameters as atkParam
 import Core.LabelManager as LabelManager
 import Core.LabelManager as LabelManager
 import Core.Statistics as Statistics
 import Core.Statistics as Statistics
 import ID2TLib.Label as Label
 import ID2TLib.Label as Label
 import ID2TLib.PcapFile as PcapFile
 import ID2TLib.PcapFile as PcapFile
+import ID2TLib.Utility as Util
 
 
 
 
 class AttackController:
 class AttackController:
@@ -42,6 +44,55 @@ class AttackController:
         :param seed: random seed for param generation
         :param seed: random seed for param generation
         :return: None
         :return: None
         """
         """
+
+        def choose_attack(input_name):
+            """"
+            Finds the attack best matching to input_name
+
+            :param input_name: The name of the attack the user put in
+            :return: The best matching attack in case one was found
+            """
+
+            import pkgutil
+            import Attack
+
+            # Find all attacks, exclude some classes
+            package = Attack
+            available_attacks = []
+            for _, name, __ in pkgutil.iter_modules(package.__path__):
+                if name != 'BaseAttack' and name != 'AttackParameters':
+                    available_attacks.append(name)
+
+            input_name = input_name.lower()
+            highest_sim = 0.0
+            highest_sim_attack = None
+            for attack in available_attacks:
+                # Compares input with one of the available attacks
+                # Makes comparison with lowercase version with generic 'attack' and 'exploit' ending removed
+                similarity = difflib.SequenceMatcher(None, input_name,
+                                                     Util.rchop(attack.lower(), ('attack', 'exploit')))\
+                    .ratio()
+                # Exact match, return appropriate attack name
+                if similarity == 1.0:
+                    return attack
+                # Found more likely match
+                if similarity > highest_sim:
+                    highest_sim = similarity
+                    highest_sim_attack = attack
+
+            # Found no exactly matching attack name, print best match and exit
+            if highest_sim >= 0.6:
+                print('Found no attack of name ' + input_name + '. The closest match was ' + highest_sim_attack +
+                      '.  Use ./id2t -l for a list of available attacks.')
+                exit(1)
+            # Found no reasonably matching attack name, recommend -l and exit
+            else:
+                print('Found no attack of name ' + input_name + ' or one similar to it.'
+                      ' Use ./id2t -l for an overview of available attacks.')
+                exit(1)
+
+        attack_name = choose_attack(attack_name)
+
         print("\nCreating attack instance of \033[1m" + attack_name + "\033[0m")
         print("\nCreating attack instance of \033[1m" + attack_name + "\033[0m")
         # Load attack class
         # Load attack class
         attack_module = importlib.import_module("Attack." + attack_name)
         attack_module = importlib.import_module("Attack." + attack_name)

+ 14 - 0
code/ID2TLib/Utility.py

@@ -352,3 +352,17 @@ def get_attacker_config(ip_source_list, ip_address: str):
         attacker_ttl_mapping[ip_address] = ttl
         attacker_ttl_mapping[ip_address] = ttl
     # return port and TTL
     # return port and TTL
     return next_port, ttl
     return next_port, ttl
+
+
+def rchop(string, endings):
+    """"
+    Returns the input string with it's ending cut off, in case it was part of 'endings'
+
+    :param string: Input string
+    :param endings: List of possible endings to be cut off
+    :return: Input string with ending cut off
+    """
+    for end in endings:
+        if string.endswith(end):
+            return string[:-len(end)]
+    return string