Browse Source

added function to remove string endings to Utility.py, improved handling to allow the omission of generic endings like 'attack' or 'exploit'

Jonathan Speth 6 years ago
parent
commit
df0abb5ea6
2 changed files with 34 additions and 9 deletions
  1. 20 9
      code/Core/AttackController.py
  2. 14 0
      code/ID2TLib/Utility.py

+ 20 - 9
code/Core/AttackController.py

@@ -7,6 +7,7 @@ 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:
@@ -45,16 +46,27 @@ class AttackController:
         """
 
         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
+            """
+
             # TODO: get list dynamically from module names
             list_of_attacks = ('DDoSAttack', 'EternalBlueExploit', 'FTPWinaXeExploit', 'JoomlaRegPrivExploit',
                                'MS17ScanAttack', 'PortscanAttack', 'SalityBotnet', 'SMBLorisAttack', 'SMBScanAttack',
                                'SQLiAttack')
 
             input_name = input_name.lower()
-            highest_sim = 0
-            highest_sim_attack = ''
+            highest_sim = 0.0
+            highest_sim_attack = None
             for attack in list_of_attacks:
-                similarity = difflib.SequenceMatcher(None, input_name, attack.lower()).ratio()
+                # 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
@@ -63,18 +75,17 @@ class AttackController:
                     highest_sim = similarity
                     highest_sim_attack = attack
 
-            # Found no exactly matching attack name, print highest match
+            # 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.')
+                      '.  Use ./id2t -l for a list of available attacks.')
                 exit(1)
-            # Found no reasonably matching attack name
+            # 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 a list of available attacks.')
+                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")

+ 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
     # 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