Controller.py 5.7 KB

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