Controller.py 5.9 KB

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