3
0

AttackController.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122
  1. import importlib
  2. import sys
  3. #Aidmar
  4. import time
  5. from Attack.AttackParameters import Parameter
  6. from ID2TLib import LabelManager
  7. from ID2TLib import Statistics
  8. from ID2TLib.Label import Label
  9. from ID2TLib.PcapFile import PcapFile
  10. class AttackController:
  11. def __init__(self, pcap_file: PcapFile, statistics_class: Statistics, label_manager: LabelManager):
  12. """
  13. Creates a new AttackController. The controller manages the attack injection, including the PCAP writing.
  14. :param statistics_class:
  15. """
  16. self.statistics = statistics_class
  17. self.pcap_file = pcap_file
  18. self.label_mgr = label_manager
  19. self.current_attack = None
  20. self.added_attacks = []
  21. # The PCAP where the attack should be injected into
  22. self.base_pcap = self.statistics.pcap_filepath
  23. def create_attack(self, attack_name: str):
  24. """
  25. Creates dynamically a new class instance based on the given attack_name.
  26. :param attack_name: The name of the attack, must correspond to the attack's class name.
  27. :return: None
  28. """
  29. print("\nCreating attack instance of \033[1m" + attack_name + "\033[0m")
  30. # Load attack class
  31. attack_module = importlib.import_module("Attack." + attack_name)
  32. attack_class = getattr(attack_module, attack_name)
  33. # Set current attack
  34. self.current_attack = attack_class(self.statistics, self.base_pcap)
  35. self.added_attacks.append(self.current_attack)
  36. def process_attack(self, attack: str, params: str):
  37. """
  38. Takes as input the name of an attack (classname) and the attack parameters as string. Parses the string of
  39. attack parameters, creates the attack by writing the attack packets and returns the path of the written pcap.
  40. :param attack: The classname of the attack to injecect.
  41. :param params: The parameters for attack customization, see attack class for supported params.
  42. :return: The file path to the created pcap file.
  43. """
  44. self.create_attack(attack)
  45. # Add attack parameters if provided
  46. print("Validating and adding attack parameters.")
  47. params_dict = []
  48. if params is not None:
  49. # Convert attack param list into dictionary
  50. for entry in params:
  51. params_dict.append(entry.split('='))
  52. params_dict = dict(params_dict)
  53. # Check if Parameter.INJECT_AT_TIMESTAMP and Parameter.INJECT_AFTER_PACKET are provided at the same time
  54. # if TRUE: delete Paramter.INJECT_AT_TIMESTAMP (lower priority) and use Parameter.INJECT_AFTER_PACKET
  55. if (Parameter.INJECT_AFTER_PACKET.value in params_dict) and (
  56. Parameter.INJECT_AT_TIMESTAMP.value in params_dict):
  57. print("CONFLICT: Parameters", Parameter.INJECT_AT_TIMESTAMP.value, "and",
  58. Parameter.INJECT_AFTER_PACKET.value,
  59. "given at the same time. Ignoring", Parameter.INJECT_AT_TIMESTAMP.value, "and using",
  60. Parameter.INJECT_AFTER_PACKET.value, "instead to derive the timestamp.")
  61. del params_dict[Parameter.INJECT_AT_TIMESTAMP.value]
  62. # Extract attack_note parameter, if not provided returns an empty string
  63. key_attack_note = "attack.note"
  64. attack_note = params_dict.get(key_attack_note, "")
  65. params_dict.pop(key_attack_note, None) # delete entry if found, otherwise return an empty string
  66. # Pass paramters to attack controller
  67. self.set_params(params_dict)
  68. else:
  69. attack_note = ""
  70. # Aidmar
  71. start_time = time.clock()
  72. # Write attack into pcap file
  73. print("Generating attack packets...", end=" ")
  74. sys.stdout.flush() # force python to print text immediately
  75. total_packets, temp_attack_pcap_path = self.current_attack.generate_attack_pcap()
  76. print("done. (total: " + str(total_packets) + " pkts.)")
  77. # Aidmar
  78. end_time = time.clock()
  79. print("Generated Attack in " + str(end_time-start_time)[:4] + " sec")
  80. # Store label into LabelManager
  81. l = Label(attack, self.get_attack_start_utime(),
  82. self.get_attack_end_utime(), attack_note)
  83. self.label_mgr.add_labels(l)
  84. return temp_attack_pcap_path
  85. def get_attack_start_utime(self):
  86. """
  87. :return: The start time (timestamp of first packet) of the attack as unix timestamp.
  88. """
  89. return self.current_attack.attack_start_utime
  90. def get_attack_end_utime(self):
  91. """
  92. :return: The end time (timestamp of last packet) of the attack as unix timestamp.
  93. """
  94. return self.current_attack.attack_end_utime
  95. def set_params(self, params: dict):
  96. """
  97. Sets the attack's parameters.
  98. :param params: The parameters in a dictionary: {parameter_name: parameter_value}
  99. :return: None
  100. """
  101. for param_key, param_value in params.items():
  102. self.current_attack.add_param_value(param_key, param_value)