3
0

Controller.py 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145
  1. import os
  2. import sys
  3. from ID2TLib.AttackController import AttackController
  4. from ID2TLib.LabelManager import LabelManager
  5. from ID2TLib.PcapFile import PcapFile
  6. from ID2TLib.Statistics import Statistics
  7. class Controller:
  8. def __init__(self, pcap_file_path: str, do_extra_tests: bool):
  9. """
  10. Creates a new Controller, acting as a central coordinator for the whole application.
  11. :param pcap_file_path:
  12. """
  13. # Fields
  14. self.pcap_src_path = pcap_file_path.strip()
  15. self.pcap_dest_path = ''
  16. self.written_pcaps = []
  17. self.do_extra_tests = do_extra_tests
  18. self.seed = None
  19. # Initialize class instances
  20. print("Input file: %s" % self.pcap_src_path)
  21. self.pcap_file = PcapFile(self.pcap_src_path)
  22. self.label_manager = LabelManager(self.pcap_src_path)
  23. self.statistics = Statistics(self.pcap_file)
  24. self.statistics.do_extra_tests = self.do_extra_tests
  25. self.statisticsDB = self.statistics.get_statistics_database()
  26. self.attack_controller = AttackController(self.pcap_file, self.statistics, self.label_manager)
  27. def load_pcap_statistics(self, flag_write_file: bool, flag_recalculate_stats: bool, flag_print_statistics: bool):
  28. """
  29. Loads the PCAP statistics either from the database, if the statistics were calculated earlier, or calculates
  30. the statistics and creates a new database.
  31. :param flag_write_file: Writes the statistics to a file.
  32. :param flag_recalculate_stats: Forces the recalculation of statistics.
  33. :param flag_print_statistics: Prints the statistics on the terminal.
  34. :return: None
  35. """
  36. self.statistics.load_pcap_statistics(flag_write_file, flag_recalculate_stats, flag_print_statistics)
  37. def process_attacks(self, attacks_config: list, seeds=[]):
  38. """
  39. Creates the attack based on the attack name and the attack parameters given in the attacks_config. The
  40. attacks_config is a list of attacks, e.g.
  41. [['PortscanAttack', 'ip.src="192.168.178.2",'dst.port=80'],['PortscanAttack', 'ip.src="10.10.10.2"]].
  42. Merges the individual temporary attack pcaps into one single pcap and merges this single pcap with the
  43. input dataset.
  44. :param attacks_config: A list of attacks with their attack parameters.
  45. """
  46. # load attacks sequentially
  47. i = 0
  48. for attack in attacks_config:
  49. if len(seeds) > i:
  50. self.attack_controller.set_seed(seed=seeds[i][0])
  51. temp_attack_pcap = self.attack_controller.process_attack(attack[0], attack[1:])
  52. self.written_pcaps.append(temp_attack_pcap)
  53. i += 1
  54. # merge attack pcaps to get single attack pcap
  55. if len(self.written_pcaps) > 1:
  56. print("\nMerging temporary attack pcaps into single pcap file...", end=" ")
  57. sys.stdout.flush() # force python to print text immediately
  58. for i in range(0, len(self.written_pcaps) - 1):
  59. attacks_pcap = PcapFile(self.written_pcaps[i])
  60. attacks_pcap_path = attacks_pcap.merge_attack(self.written_pcaps[i + 1])
  61. os.remove(self.written_pcaps[i + 1]) # remove merged pcap
  62. self.written_pcaps[i + 1] = attacks_pcap_path
  63. print("done.")
  64. else:
  65. attacks_pcap_path = self.written_pcaps[0]
  66. # merge single attack pcap with all attacks into base pcap
  67. print("Merging base pcap with single attack pcap...", end=" ")
  68. sys.stdout.flush() # force python to print text immediately
  69. self.pcap_dest_path = self.pcap_file.merge_attack(attacks_pcap_path)
  70. tmp_path_tuple = self.pcap_dest_path.rpartition("/")
  71. result_dir = tmp_path_tuple[0] + tmp_path_tuple[1] + "ID2T_results/"
  72. result_path = result_dir + tmp_path_tuple[2]
  73. os.makedirs(result_dir, exist_ok=True)
  74. os.rename(self.pcap_dest_path, result_path)
  75. self.pcap_dest_path = result_path
  76. print("done.")
  77. # delete intermediate PCAP files
  78. print('Deleting intermediate attack pcap...', end=" ")
  79. sys.stdout.flush() # force python to print text immediately
  80. os.remove(attacks_pcap_path)
  81. print("done.")
  82. # write label file with attacks
  83. self.label_manager.write_label_file(self.pcap_dest_path)
  84. # print status message
  85. print('\nOutput files created: \n', self.pcap_dest_path, '\n', self.label_manager.label_file_path)
  86. def process_db_queries(self, query, print_results=False):
  87. """
  88. Processes a statistics database query. This can be a standard SQL query or a named query.
  89. :param query: The query as a string or multiple queries as a list of strings.
  90. :param print_results: Must be True if the results should be printed to terminal.
  91. :return: The query's result
  92. """
  93. print("Processing database query/queries...")
  94. if isinstance(query, list) or isinstance(query, tuple):
  95. for q in query:
  96. self.statisticsDB.process_db_query(q, print_results)
  97. else:
  98. self.statisticsDB.process_db_query(query, print_results)
  99. def enter_query_mode(self):
  100. """
  101. Enters into the query mode. This is a read-eval-print-loop, where the user can input named queries or SQL
  102. queries and the results are printed.
  103. """
  104. print("Entering into query mode...")
  105. print("Enter statement ending by ';' and press ENTER to send query. Exit by sending an empty query..")
  106. buffer = ""
  107. while True:
  108. line = input("> ")
  109. if line == "":
  110. break
  111. buffer += line
  112. import sqlite3
  113. if sqlite3.complete_statement(buffer):
  114. try:
  115. buffer = buffer.strip()
  116. self.statisticsDB.process_db_query(buffer, True)
  117. except sqlite3.Error as e:
  118. print("An error occurred:", e.args[0])
  119. buffer = ""
  120. def create_statistics_plot(self, params: str):
  121. """
  122. Plots the statistics to a file by using the given customization parameters.
  123. """
  124. if params is not None and params[0] is not None:
  125. params_dict = dict([z.split("=") for z in params])
  126. self.statistics.plot_statistics(format=params_dict['format'])
  127. else:
  128. self.statistics.plot_statistics()