Merge branch 'attack_recognition' of stefan.schmidt/ID2T-toolkit into master

code/Core/AttackController.py

@@ -1,11 +1,13 @@
 import importlib
 import sys
+import difflib
 
 import Attack.AttackParameters as atkParam
 import Core.LabelManager as LabelManager
 import Core.Statistics as Statistics
 import ID2TLib.Label as Label
 import ID2TLib.PcapFile as PcapFile
+import ID2TLib.Utility as Util
 
 
 class AttackController:
@@ -42,6 +44,55 @@ class AttackController:
         :param seed: random seed for param generation
         :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")
         # Load attack class
         attack_module = importlib.import_module("Attack." + attack_name)

code/ID2TLib/Utility.py

@@ -352,3 +352,17 @@ def get_attacker_config(ip_source_list, ip_address: str):
         attacker_ttl_mapping[ip_address] = ttl
     # return port and 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