Controller.py 8.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186
  1. import os
  2. import sys
  3. import shutil
  4. import time
  5. from ID2TLib.AttackController import AttackController
  6. from ID2TLib.LabelManager import LabelManager
  7. from ID2TLib.PcapFile import PcapFile
  8. from ID2TLib.Statistics import Statistics
  9. from ID2TLib.AttackContext import AttackContext
  10. class Controller:
  11. def __init__(self, in_pcap_file_path: str, do_extra_tests: bool, out_pcap_file_path):
  12. """
  13. Creates a new Controller, acting as a central coordinator for the whole application.
  14. :param pcap_file_path:
  15. """
  16. # Fields
  17. self.pcap_src_path = in_pcap_file_path.strip()
  18. if out_pcap_file_path:
  19. self.pcap_out_path = out_pcap_file_path.strip()
  20. else:
  21. self.pcap_out_path = None
  22. self.pcap_dest_path = ''
  23. self.written_pcaps = []
  24. self.do_extra_tests = do_extra_tests
  25. # Initialize class instances
  26. print("Input file: %s" % self.pcap_src_path)
  27. self.pcap_file = PcapFile(self.pcap_src_path)
  28. self.label_manager = LabelManager(self.pcap_src_path)
  29. self.statistics = Statistics(self.pcap_file)
  30. self.statistics.do_extra_tests = self.do_extra_tests
  31. self.statisticsDB = self.statistics.get_statistics_database()
  32. self.attack_controller = AttackController(self.pcap_file, self.statistics, self.label_manager)
  33. def load_pcap_statistics(self, flag_write_file: bool, flag_recalculate_stats: bool, flag_print_statistics: bool):
  34. """
  35. Loads the PCAP statistics either from the database, if the statistics were calculated earlier, or calculates
  36. the statistics and creates a new database.
  37. :param flag_write_file: Writes the statistics to a file.
  38. :param flag_recalculate_stats: Forces the recalculation of statistics.
  39. :param flag_print_statistics: Prints the statistics on the terminal.
  40. :return: None
  41. """
  42. self.statistics.load_pcap_statistics(flag_write_file, flag_recalculate_stats, flag_print_statistics)
  43. def process_attacks(self, attacks_config: list, inject_empty: bool=False):
  44. """
  45. Creates the attack based on the attack name and the attack parameters given in the attacks_config. The
  46. attacks_config is a list of attacks, e.g.
  47. [['PortscanAttack', 'ip.src="192.168.178.2",'dst.port=80'],['PortscanAttack', 'ip.src="10.10.10.2"]].
  48. Merges the individual temporary attack pcaps into one single pcap and merges this single pcap with the
  49. input dataset if desired.
  50. :param attacks_config: A list of attacks with their attack parameters.
  51. :param inject_empty: if flag is set, Attack PCAPs will not be merged with the base PCAP, ie. Attacks are injected into an empty PCAP
  52. """
  53. # get output directory
  54. if self.pcap_out_path:
  55. out_dir = os.path.dirname(self.pcap_out_path)
  56. else:
  57. out_dir = os.path.dirname(self.pcap_src_path)
  58. # if out_dir is cwd
  59. if out_dir == "":
  60. out_dir = "."
  61. # context for the attack(s)
  62. context = AttackContext(out_dir)
  63. # note if new xml file has been created by MembersMgmtCommAttack
  64. # load attacks sequentially
  65. for attack in attacks_config:
  66. temp_attack_pcap = self.attack_controller.process_attack(attack[0], attack[1:], context)
  67. self.written_pcaps.append(temp_attack_pcap)
  68. # merge attack pcaps to get single attack pcap
  69. if len(self.written_pcaps) > 1:
  70. print("\nMerging temporary attack pcaps into single pcap file...", end=" ")
  71. sys.stdout.flush() # force python to print text immediately
  72. for i in range(0, len(self.written_pcaps) - 1):
  73. attacks_pcap = PcapFile(self.written_pcaps[i])
  74. attacks_pcap_path = attacks_pcap.merge_attack(self.written_pcaps[i + 1])
  75. os.remove(self.written_pcaps[i + 1]) # remove merged pcap
  76. print("done.")
  77. else:
  78. attacks_pcap_path = self.written_pcaps[0]
  79. if inject_empty:
  80. # copy the attack pcap to the directory of the base PCAP instead of merging them
  81. print("Copying single attack pcap to location of base pcap...", end=" ")
  82. sys.stdout.flush() # force python to print text immediately
  83. timestamp = time.strftime("%Y%m%d") + '-' + time.strftime("%X").replace(':', '')
  84. self.pcap_dest_path = self.pcap_src_path.replace(".pcap", timestamp + '.pcap')
  85. shutil.copy(attacks_pcap_path, self.pcap_dest_path)
  86. else:
  87. # merge single attack pcap with all attacks into base pcap
  88. print("Merging base pcap with single attack pcap...", end=" ")
  89. sys.stdout.flush() # force python to print text immediately
  90. # cp merged PCAP to output path
  91. self.pcap_dest_path = self.pcap_file.merge_attack(attacks_pcap_path)
  92. if self.pcap_out_path:
  93. if not self.pcap_out_path.endswith(".pcap"):
  94. self.pcap_out_path += ".pcap"
  95. os.rename(self.pcap_dest_path, self.pcap_out_path)
  96. self.pcap_dest_path = self.pcap_out_path
  97. print("done.")
  98. # delete intermediate PCAP files
  99. print('Deleting intermediate attack pcap...', end=" ")
  100. sys.stdout.flush() # force python to print text immediately
  101. os.remove(attacks_pcap_path)
  102. print("done.")
  103. # write label file with attacks
  104. self.label_manager.write_label_file(self.pcap_dest_path)
  105. # pcap_base contains the name of the pcap-file without the ".pcap" extension
  106. pcap_base = os.path.splitext(self.pcap_dest_path)[0]
  107. created_files = [self.pcap_dest_path, self.label_manager.label_file_path]
  108. for suffix, filename in context.get_allocated_files():
  109. shutil.move(filename, pcap_base + suffix)
  110. created_files.append(pcap_base + suffix)
  111. context.reset()
  112. # print status message
  113. created_files += context.get_other_created_files()
  114. created_files.sort()
  115. print("\nOutput files created:")
  116. for file in created_files:
  117. # remove ./ at beginning of file to have only one representation for cwd
  118. if file.startswith("./"):
  119. file = file[2:]
  120. print(file)
  121. def process_db_queries(self, query, print_results=False):
  122. """
  123. Processes a statistics database query. This can be a standard SQL query or a named query.
  124. :param query: The query as a string or multiple queries as a list of strings.
  125. :param print_results: Must be True if the results should be printed to terminal.
  126. :return: The query's result
  127. """
  128. print("Processing database query/queries...")
  129. if isinstance(query, list) or isinstance(query, tuple):
  130. for q in query:
  131. self.statisticsDB.process_db_query(q, print_results)
  132. else:
  133. self.statisticsDB.process_db_query(query, print_results)
  134. def enter_query_mode(self):
  135. """
  136. Enters into the query mode. This is a read-eval-print-loop, where the user can input named queries or SQL
  137. queries and the results are printed.
  138. """
  139. print("Entering into query mode...")
  140. print("Enter statement ending by ';' and press ENTER to send query. Exit by sending an empty query..")
  141. buffer = ""
  142. while True:
  143. line = input("> ")
  144. if line == "":
  145. break
  146. buffer += line
  147. import sqlite3
  148. if sqlite3.complete_statement(buffer):
  149. try:
  150. buffer = buffer.strip()
  151. self.statisticsDB.process_db_query(buffer, True)
  152. except sqlite3.Error as e:
  153. print("An error occurred:", e.args[0])
  154. buffer = ""
  155. def create_statistics_plot(self, params: str):
  156. """
  157. Plots the statistics to a file by using the given customization parameters.
  158. """
  159. if params is not None and params[0] is not None:
  160. params_dict = dict([z.split("=") for z in params])
  161. self.statistics.plot_statistics(format=params_dict['format'])
  162. else:
  163. self.statistics.plot_statistics()