Controller.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275
  1. import os
  2. import readline
  3. import sys
  4. import ID2TLib.AttackController as atkCtrl
  5. import ID2TLib.LabelManager as LabelManager
  6. import ID2TLib.PcapFile as PcapFile
  7. import ID2TLib.Statistics as Statistics
  8. class Controller:
  9. def __init__(self, pcap_file_path: str, do_extra_tests: bool):
  10. """
  11. Creates a new Controller, acting as a central coordinator for the whole application.
  12. :param pcap_file_path:
  13. """
  14. # Fields
  15. self.pcap_src_path = pcap_file_path.strip()
  16. self.pcap_dest_path = ''
  17. self.written_pcaps = []
  18. self.do_extra_tests = do_extra_tests
  19. self.seed = None
  20. self.durations = []
  21. # Initialize class instances
  22. print("Input file: %s" % self.pcap_src_path)
  23. self.pcap_file = PcapFile.PcapFile(self.pcap_src_path)
  24. self.label_manager = LabelManager.LabelManager(self.pcap_src_path)
  25. self.statistics = Statistics.Statistics(self.pcap_file)
  26. self.statistics.do_extra_tests = self.do_extra_tests
  27. self.statisticsDB = self.statistics.get_statistics_database()
  28. self.attack_controller = atkCtrl.AttackController(self.pcap_file, self.statistics, self.label_manager)
  29. def load_pcap_statistics(self, flag_write_file: bool, flag_recalculate_stats: bool, flag_print_statistics: bool):
  30. """
  31. Loads the PCAP statistics either from the database, if the statistics were calculated earlier, or calculates
  32. the statistics and creates a new database.
  33. :param flag_write_file: Writes the statistics to a file.
  34. :param flag_recalculate_stats: Forces the recalculation of statistics.
  35. :param flag_print_statistics: Prints the statistics on the terminal.
  36. :return: None
  37. """
  38. self.statistics.load_pcap_statistics(flag_write_file, flag_recalculate_stats, flag_print_statistics)
  39. def process_attacks(self, attacks_config: list, seeds=None, time=False):
  40. """
  41. Creates the attack based on the attack name and the attack parameters given in the attacks_config. The
  42. attacks_config is a list of attacks.
  43. e.g. [['PortscanAttack', 'ip.src="192.168.178.2",'dst.port=80'],['PortscanAttack', 'ip.src="10.10.10.2"]].
  44. Merges the individual temporary attack pcaps into one single pcap and merges this single pcap with the
  45. input dataset.
  46. :param attacks_config: A list of attacks with their attack parameters.
  47. :param seeds: A list of random seeds for the given attacks.
  48. :param time: Measure time for packet generation.
  49. """
  50. # load attacks sequentially
  51. i = 0
  52. for attack in attacks_config:
  53. if seeds is not None and len(seeds) > i:
  54. self.attack_controller.set_seed(seed=seeds[i][0])
  55. temp_attack_pcap, duration = self.attack_controller.process_attack(attack[0], attack[1:], time)
  56. self.durations.append(duration)
  57. self.written_pcaps.append(temp_attack_pcap)
  58. i += 1
  59. attacks_pcap_path = None
  60. # merge attack pcaps to get single attack pcap
  61. if len(self.written_pcaps) > 1:
  62. print("\nMerging temporary attack pcaps into single pcap file...", end=" ")
  63. sys.stdout.flush() # force python to print text immediately
  64. for i in range(0, len(self.written_pcaps) - 1):
  65. attacks_pcap = PcapFile.PcapFile(self.written_pcaps[i])
  66. attacks_pcap_path = attacks_pcap.merge_attack(self.written_pcaps[i + 1])
  67. os.remove(self.written_pcaps[i + 1]) # remove merged pcap
  68. self.written_pcaps[i + 1] = attacks_pcap_path
  69. print("done.")
  70. else:
  71. attacks_pcap_path = self.written_pcaps[0]
  72. # merge single attack pcap with all attacks into base pcap
  73. print("Merging base pcap with single attack pcap...", end=" ")
  74. sys.stdout.flush() # force python to print text immediately
  75. self.pcap_dest_path = self.pcap_file.merge_attack(attacks_pcap_path)
  76. tmp_path_tuple = self.pcap_dest_path.rpartition("/")
  77. result_dir = tmp_path_tuple[0] + tmp_path_tuple[1] + "ID2T_results/"
  78. result_path = result_dir + tmp_path_tuple[2]
  79. os.makedirs(result_dir, exist_ok=True)
  80. os.rename(self.pcap_dest_path, result_path)
  81. self.pcap_dest_path = result_path
  82. print("done.")
  83. # delete intermediate PCAP files
  84. print('Deleting intermediate attack pcap...', end=" ")
  85. sys.stdout.flush() # force python to print text immediately
  86. os.remove(attacks_pcap_path)
  87. print("done.")
  88. # write label file with attacks
  89. self.label_manager.write_label_file(self.pcap_dest_path)
  90. # print status message
  91. print('\nOutput files created: \n', self.pcap_dest_path, '\n', self.label_manager.label_file_path)
  92. def process_db_queries(self, query, print_results=False):
  93. """
  94. Processes a statistics database query. This can be a standard SQL query or a named query.
  95. :param query: The query as a string or multiple queries as a list of strings.
  96. :param print_results: Must be True if the results should be printed to terminal.
  97. :return: The query's result
  98. """
  99. print("Processing database query/queries...")
  100. if isinstance(query, list) or isinstance(query, tuple):
  101. for q in query:
  102. self.statisticsDB.process_db_query(q, print_results)
  103. else:
  104. self.statisticsDB.process_db_query(query, print_results)
  105. @staticmethod
  106. def process_help(params):
  107. if not params:
  108. print("Query mode allows you to enter SQL-queries as well as named queries.")
  109. print()
  110. print("Named queries:")
  111. print("\tSelectors:")
  112. print("\t\tmost_used(...) -> Returns the most occurring element in all elements")
  113. print("\t\tleast_used(...) -> Returns the least occurring element in all elements")
  114. print("\t\tavg(...) -> Returns the average of all elements")
  115. print("\t\tall(...) -> Returns all elements")
  116. print("\tExtractors:")
  117. print("\t\trandom(...) -> Returns a random element from a list")
  118. print("\t\tfirst(...) -> Returns the first element from a list")
  119. print("\t\tlast(...) -> Returns the last element from a list")
  120. print("\tParameterized selectors:")
  121. print("\t\tipAddress(...) -> Returns all IP addresses fulfilling the specified conditions")
  122. print("\t\tmacAddress(...) -> Returns all MAC addresses fulfilling the specified conditions")
  123. print()
  124. print("Miscellaneous:")
  125. print("\tlabels -> List all attacks listed in the label file, if any")
  126. print()
  127. print("Additional information is available with 'help [KEYWORD];'")
  128. print("To get a list of examples, type 'help examples;'")
  129. print()
  130. return
  131. param = params[0].lower()
  132. if param == "most_used":
  133. print("most_used can be used as a selector for the following attributes:")
  134. print("ipAddress | macAddress | portNumber | protocolName | ttlValue | mssValue | winSize | ipClass")
  135. print()
  136. elif param == "least_used":
  137. print("least_used can be used as a selector for the following attributes:")
  138. print("ipAddress | macAddress | portNumber | protocolName | ttlValue")
  139. print()
  140. elif param == "avg":
  141. print("avg can be used as a selector for the following attributes:")
  142. print("pktsReceived | pktsSent | kbytesSent | kbytesReceived | ttlValue | mss")
  143. print()
  144. elif param == "all":
  145. print("all can be used as a selector for the following attributes:")
  146. print("ipAddress | ttlValue | mss | macAddress | portNumber | protocolName")
  147. print()
  148. elif param in ["random", "first", "last"]:
  149. print("No additional info available for this keyword.")
  150. print()
  151. elif param == "ipaddress":
  152. print("ipAddress is a parameterized selector which fetches IP addresses based on (a list of) conditions.")
  153. print("Conditions are of the following form: PARAMETER OPERATOR VALUE")
  154. print("The following parameters can be specified:")
  155. print("pktsReceived | pktsSent | kbytesReceived | kbytesSent | maxPktRate | minPktRate | ipClass\n"
  156. "macAddress | ttlValue | ttlCount | portDirection | portNumber | portCount | protocolCount\n"
  157. "protocolName")
  158. print()
  159. print("See 'help examples;' for usage examples.")
  160. print()
  161. elif param == "macaddress":
  162. print("macAddress is a parameterized selector which fetches MAC addresses based on (a list of) conditions.")
  163. print("Conditions are of the following form: PARAMETER OPERATOR VALUE")
  164. print("The following parameters can be specified:")
  165. print("ipAddress")
  166. print()
  167. print("See 'help examples;' for usage examples.")
  168. print()
  169. elif param == "examples":
  170. print("Get the average amount of sent packets per IP:")
  171. print("\tavg(pktsSent);")
  172. print("Get a random IP from all addresses occuring in the pcap:")
  173. print("\trandom(all(ipAddress));")
  174. print("Return the MAC address of a specified IP:")
  175. print("\tmacAddress(ipAddress=192.168.178.2);")
  176. print("Get the average TTL-value with SQL:")
  177. print("\tSELECT avg(ttlValue) from ip_ttl;")
  178. print("Get a random IP address from all addresses that sent and received at least 10 packets:")
  179. print("\trandom(ipAddress(pktsSent > 10, pktsReceived > 10));")
  180. print()
  181. else:
  182. print("Unknown keyword '" + param + "', try 'help;' to get a list of allowed keywords'")
  183. print()
  184. def enter_query_mode(self):
  185. """
  186. Enters into the query mode. This is a read-eval-print-loop, where the user can input named queries or SQL
  187. queries and the results are printed.
  188. """
  189. def make_completer(vocabulary):
  190. def custom_template(text, state):
  191. results = [x for x in vocabulary if x.startswith(text)] + [None]
  192. return results[state]
  193. return custom_template
  194. readline.parse_and_bind('tab: complete')
  195. readline.set_completer(make_completer(
  196. self.statisticsDB.get_all_named_query_keywords() + self.statisticsDB.get_all_sql_query_keywords()))
  197. history_file = os.path.join(os.path.expanduser('~'), 'ID2T_data', 'query_history')
  198. try:
  199. readline.read_history_file(history_file)
  200. except IOError:
  201. pass
  202. print("Entering into query mode...")
  203. print("Enter statement ending by ';' and press ENTER to send query. Exit by sending an empty query.")
  204. print("Type 'help;' for information on possible queries.")
  205. buffer = ""
  206. while True:
  207. line = input("> ")
  208. if line == "":
  209. break
  210. buffer += line
  211. import sqlite3
  212. if sqlite3.complete_statement(buffer):
  213. try:
  214. buffer = buffer.strip()
  215. if buffer.lower().startswith('help'):
  216. buffer = buffer.strip(';')
  217. self.process_help(buffer.split(' ')[1:])
  218. elif buffer.lower().strip() == 'labels;':
  219. if not self.label_manager.labels:
  220. print("No labels found.")
  221. else:
  222. print("Attacks listed in the label file:")
  223. print()
  224. for label in self.label_manager.labels:
  225. print("Attack name: " + str(label.attack_name))
  226. print("Attack note: " + str(label.attack_note))
  227. print("Start timestamp: " + str(label.timestamp_start))
  228. print("End timestamp: " + str(label.timestamp_end))
  229. print()
  230. print()
  231. else:
  232. self.statisticsDB.process_db_query(buffer, True)
  233. except sqlite3.Error as e:
  234. print("An error occurred:", e.args[0])
  235. buffer = ""
  236. readline.set_history_length(1000)
  237. readline.write_history_file(history_file)
  238. def create_statistics_plot(self, params: str, entropy: bool):
  239. """
  240. Plots the statistics to a file by using the given customization parameters.
  241. """
  242. if params is not None and params[0] is not None:
  243. # FIXME: cleanup
  244. params_dict = dict([z.split("=") for z in params])
  245. self.statistics.plot_statistics(entropy=entropy, file_format=params_dict['format'])
  246. else:
  247. self.statistics.plot_statistics(entropy=entropy)