Controller.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428
  1. import os
  2. import readline
  3. import sys
  4. import shutil
  5. import time
  6. import re
  7. import pyparsing as pp
  8. import Core.AttackController as atkCtrl
  9. import Core.LabelManager as LabelManager
  10. import Core.Statistics as Statistics
  11. import ID2TLib.PcapFile as PcapFile
  12. import ID2TLib.Utility as Util
  13. import Core.StatsDatabase as StatsDB
  14. class Controller:
  15. def __init__(self, pcap_file_path: str, do_extra_tests: bool, non_verbose: bool=True, pcap_out_path: str=None,
  16. debug: bool=False):
  17. """
  18. Creates a new Controller, acting as a central coordinator for the whole application.
  19. :param pcap_file_path:
  20. """
  21. # Fields
  22. self.pcap_src_path = pcap_file_path.strip()
  23. self.pcap_dest_path = ''
  24. self.pcap_out_path = pcap_out_path
  25. self.written_pcaps = []
  26. self.do_extra_tests = do_extra_tests
  27. self.non_verbose = non_verbose
  28. self.seed = None
  29. self.durations = []
  30. self.added_packets = 0
  31. self.created_files = []
  32. self.debug = debug
  33. # Initialize class instances
  34. print("Input file: %s" % self.pcap_src_path)
  35. self.pcap_file = PcapFile.PcapFile(self.pcap_src_path)
  36. self.label_manager = LabelManager.LabelManager(self.pcap_src_path)
  37. self.statistics = Statistics.Statistics(self.pcap_file)
  38. self.statistics.do_extra_tests = self.do_extra_tests
  39. self.statisticsDB = self.statistics.get_statistics_database()
  40. self.attack_controller = atkCtrl.AttackController(self.pcap_file, self.statistics, self.label_manager)
  41. # Set output directory and create it (if necessary)
  42. if pcap_out_path is not None:
  43. out_dir = os.path.dirname(pcap_out_path)
  44. if not out_dir: # if out_dir is cwd
  45. out_dir = "."
  46. Util.OUT_DIR = out_dir + os.sep
  47. else:
  48. Util.OUT_DIR = os.path.join(os.path.dirname(pcap_file_path), "ID2T_results") + os.sep
  49. os.makedirs(Util.OUT_DIR, exist_ok=True)
  50. def list_interval_statistics(self):
  51. self.statistics.list_previous_interval_statistic_tables()
  52. def load_pcap_statistics(self, flag_write_file: bool, flag_recalculate_stats: bool, flag_print_statistics: bool,
  53. intervals, delete: bool=False, recalculate_intervals: bool=None):
  54. """
  55. Loads the PCAP statistics either from the database, if the statistics were calculated earlier, or calculates
  56. the statistics and creates a new database.
  57. :param flag_write_file: Writes the statistics to a file.
  58. :param flag_recalculate_stats: Forces the recalculation of statistics.
  59. :param flag_print_statistics: Prints the statistics on the terminal.
  60. :param intervals: user specified interval in seconds
  61. :param delete: Delete old interval statistics.
  62. :param recalculate_intervals: Recalculate old interval statistics or not. Prompt user if None.
  63. :return: None
  64. """
  65. self.statistics.load_pcap_statistics(flag_write_file, flag_recalculate_stats, flag_print_statistics,
  66. self.non_verbose, intervals=intervals, delete=delete,
  67. recalculate_intervals=recalculate_intervals)
  68. def process_attacks(self, attacks_config: list, seeds=None, time: bool=False, inject_empty: bool=False):
  69. """
  70. Creates the attack based on the attack name and the attack parameters given in the attacks_config. The
  71. attacks_config is a list of attacks.
  72. e.g. [['PortscanAttack', 'ip.src="192.168.178.2",'dst.port=80'],['PortscanAttack', 'ip.src="10.10.10.2"]].
  73. Merges the individual temporary attack pcaps into one single pcap and merges this single pcap with the
  74. input dataset if desired.
  75. :param attacks_config: A list of attacks with their attack parameters.
  76. :param seeds: A list of random seeds for the given attacks.
  77. :param time: Measure time for packet generation.
  78. :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
  79. """
  80. # load attacks sequentially
  81. i = 0
  82. for attack in attacks_config:
  83. if seeds is not None and len(seeds) > i:
  84. rng_seed = seeds[i][0]
  85. else:
  86. rng_seed = int.from_bytes(os.urandom(16), sys.byteorder)
  87. self.attack_controller.set_seed(seed=rng_seed)
  88. temp_attack_pcap, duration = self.attack_controller.process_attack(attack[0], attack[1:], time)
  89. self.durations.append(duration)
  90. self.added_packets += self.attack_controller.total_packets
  91. if not self.non_verbose:
  92. self.statistics.stats_summary_post_attack(self.added_packets)
  93. self.written_pcaps.append(temp_attack_pcap)
  94. i += 1
  95. attacks_pcap_path = None
  96. # merge attack pcaps to get single attack pcap
  97. if len(self.written_pcaps) > 1:
  98. print("\nMerging temporary attack pcaps into single pcap file...", end=" ")
  99. sys.stdout.flush() # force python to print text immediately
  100. for i in range(0, len(self.written_pcaps) - 1):
  101. attacks_pcap = PcapFile.PcapFile(self.written_pcaps[i])
  102. attacks_pcap_path = attacks_pcap.merge_attack(self.written_pcaps[i + 1])
  103. os.remove(self.written_pcaps[i + 1]) # remove merged pcap
  104. self.written_pcaps[i + 1] = attacks_pcap_path
  105. print("done.")
  106. elif len(self.written_pcaps) == 1:
  107. attacks_pcap_path = self.written_pcaps[0]
  108. if attacks_pcap_path:
  109. if inject_empty:
  110. # copy the attack pcap to the directory of the base PCAP instead of merging them
  111. print("Copying single attack pcap to location of base pcap...", end=" ")
  112. sys.stdout.flush() # force python to print text immediately
  113. timestamp = '_' + time.strftime("%Y%m%d") + '-' + time.strftime("%X").replace(':', '')
  114. self.pcap_dest_path = self.pcap_src_path.replace(".pcap", timestamp + '.pcap')
  115. shutil.copy(attacks_pcap_path, self.pcap_dest_path)
  116. else:
  117. # merge single attack pcap with all attacks into base pcap
  118. print("Merging base pcap with single attack pcap...", end=" ")
  119. sys.stdout.flush() # force python to print text immediately
  120. self.pcap_dest_path = self.pcap_file.merge_attack(attacks_pcap_path)
  121. if self.pcap_out_path:
  122. if not self.pcap_out_path.endswith(".pcap"):
  123. self.pcap_out_path += ".pcap"
  124. result_path = self.pcap_out_path
  125. else:
  126. tmp_path_tuple = self.pcap_dest_path.rpartition("/")
  127. result_path = Util.OUT_DIR + tmp_path_tuple[2]
  128. os.rename(self.pcap_dest_path, result_path)
  129. self.pcap_dest_path = result_path
  130. self.created_files = [self.pcap_dest_path]
  131. # process/move other created files
  132. pcap_basename = os.path.splitext(self.pcap_dest_path)[0]
  133. for x in self.attack_controller.additional_files:
  134. outpath = pcap_basename + "_" + x
  135. os.rename(x, outpath)
  136. self.created_files.append(outpath)
  137. print("done.")
  138. # delete intermediate PCAP files
  139. if self.debug:
  140. print('NOT deleting intermediate attack pcap while in debug mode.')
  141. else:
  142. print('Deleting intermediate attack pcap...', end=" ")
  143. sys.stdout.flush() # force python to print text immediately
  144. os.remove(attacks_pcap_path)
  145. print("done.")
  146. # write label file with attacks
  147. self.label_manager.write_label_file(self.pcap_dest_path)
  148. self.created_files.insert(1, self.label_manager.label_file_path)
  149. # print status message
  150. print('\nOutput files created:')
  151. for filepath in self.created_files:
  152. print(filepath)
  153. else:
  154. print("done.")
  155. print('\nOutput files created:')
  156. print("--> No packets were injected. Therefore no output files were created.")
  157. # print summary statistics
  158. if not self.non_verbose and len(attacks_config) is not 1:
  159. self.statistics.stats_summary_post_attack(self.added_packets)
  160. def process_db_queries(self, query, print_results=False):
  161. """
  162. Processes a statistics database query. This can be a standard SQL query or a named query.
  163. :param query: The query as a string or multiple queries as a list of strings.
  164. :param print_results: Must be True if the results should be printed to terminal.
  165. :return: The query's result
  166. """
  167. print("Processing database query/queries...")
  168. if isinstance(query, list) or isinstance(query, tuple):
  169. for q in query:
  170. self.statisticsDB.process_db_query(q, print_results)
  171. else:
  172. self.statisticsDB.process_db_query(query, print_results)
  173. @staticmethod
  174. def process_help(params) -> None:
  175. """
  176. Prints either general help messages, or information about specific commands.
  177. :param params: A list of parameters for the help command (can be empty).
  178. """
  179. if not params:
  180. print("Query mode allows you to enter SQL-queries as well as named queries.")
  181. print()
  182. print("Named queries:")
  183. print("\tSelectors:")
  184. print("\t\tmost_used(...) -> Returns the most occurring element in all elements")
  185. print("\t\tleast_used(...) -> Returns the least occurring element in all elements")
  186. print("\t\tavg(...) -> Returns the average of all elements")
  187. print("\t\tall(...) -> Returns all elements")
  188. print("\tExtractors:")
  189. print("\t\trandom(...) -> Returns a random element from a list")
  190. print("\t\tfirst(...) -> Returns the first element from a list")
  191. print("\t\tlast(...) -> Returns the last element from a list")
  192. print("\tParameterized selectors:")
  193. print("\t\tipAddress(...) -> Returns all IP addresses fulfilling the specified conditions")
  194. print("\t\tmacAddress(...) -> Returns all MAC addresses fulfilling the specified conditions")
  195. print()
  196. print("Miscellaneous:")
  197. print("\tlabels -> List all attacks listed in the label file, if any")
  198. print("\ttables -> List all tables from database")
  199. print("\tcolumns TABLE -> List column names and types from specified table")
  200. print()
  201. print("Additional information is available with 'help [KEYWORD];'")
  202. print("To get a list of examples, type 'help examples;'")
  203. print()
  204. return
  205. param = params[0].lower()
  206. if param == "most_used":
  207. print("most_used can be used as a selector for the following attributes:")
  208. print("ipAddress | macAddress | portNumber | protocolName | ttlValue | mssValue | winSize | ipClass")
  209. print()
  210. elif param == "least_used":
  211. print("least_used can be used as a selector for the following attributes:")
  212. print("ipAddress | macAddress | portNumber | protocolName | ttlValue | mssValue | winSize | ipClass")
  213. print()
  214. elif param == "avg":
  215. print("avg can be used as a selector for the following attributes:")
  216. print("pktsReceived | pktsSent | kbytesSent | kbytesReceived | ttlValue | mss")
  217. print()
  218. elif param == "all":
  219. print("all can be used as a selector for the following attributes:")
  220. print("ipAddress | ttlValue | mss | macAddress | portNumber | protocolName | winSize | ipClass")
  221. print()
  222. elif param in ["random", "first", "last"]:
  223. print("No additional info available for this keyword.")
  224. print()
  225. elif param == "ipaddress":
  226. print("ipAddress is a parameterized selector which fetches IP addresses based on (a list of) conditions.")
  227. print("Conditions are of the following form: PARAMETER OPERATOR VALUE")
  228. print("The following parameters can be specified:")
  229. print("pktsReceived | pktsSent | kbytesReceived | kbytesSent | maxPktRate | minPktRate | ipClass\n"
  230. "macAddress | ttlValue | ttlCount | portDirection | portNumber | portCount | protocolCount\n"
  231. "protocolName")
  232. print()
  233. print("The following operators can be used:")
  234. print("<= | < | = | >= | > | in")
  235. print()
  236. print("A value can either be a simple values, a list of simple values separated by commas and enclosed "
  237. "in [] brackets, or another query.")
  238. print()
  239. print("When VALUE is a list (or a query returning a list), the usage of the 'in' operator is mandatory!")
  240. print()
  241. print("See 'help examples;' for usage examples.")
  242. print()
  243. elif param == "macaddress":
  244. print("macAddress is a parameterized selector which fetches MAC addresses based on (a list of) conditions.")
  245. print("Conditions are of the following form: PARAMETER OPERATOR VALUE")
  246. print("The following parameters can be specified:")
  247. print("ipAddress")
  248. print()
  249. print("See 'help ipAddress' for information on valid operators and values.")
  250. print()
  251. print("See 'help examples;' for usage examples.")
  252. print()
  253. elif param == "examples":
  254. print("Get the average amount of sent packets per IP:")
  255. print("\tavg(pktsSent);")
  256. print("Get a random IP from all addresses occuring in the pcap:")
  257. print("\trandom(all(ipAddress));")
  258. print("Return the MAC address of a specified IP:")
  259. print("\tmacAddress(ipAddress=192.168.178.2);")
  260. print("Get the average TTL-value with SQL:")
  261. print("\tSELECT avg(ttlValue) from ip_ttl;")
  262. print("Get a random IP address from all addresses that sent and received at least 10 packets:")
  263. print("\trandom(ipAddress(pktsSent > 10, pktsReceived > 10));")
  264. print("Get the IP addresses used with one of the MAC addresses in a list:")
  265. print("\tipAddress(macAddress in [08:00:27:a3:83:43, 52:54:00:12:35:02]);")
  266. print()
  267. else:
  268. print("Unknown keyword '" + param + "', try 'help;' to get a list of allowed keywords'")
  269. print()
  270. def internal_command(self, query: str) -> bool:
  271. # Strip off semicolon, split into command and parameters
  272. query = query.strip(";").split(" ", 1)
  273. cmd = query[0].strip().lower()
  274. if len(query) > 1:
  275. params = [p for p in re.split("(,|\\\".*?\\\"|'.*?')", query[1]) if p.strip(",").strip()]
  276. params = list(map(lambda x: x.strip().strip("\"'"), params))
  277. else:
  278. params = []
  279. if cmd == "help":
  280. self.process_help(params)
  281. return True
  282. elif cmd == "labels":
  283. if not self.label_manager.labels:
  284. print("No labels found.")
  285. else:
  286. print("Attacks listed in the label file:")
  287. print()
  288. for i, label in enumerate(self.label_manager.labels):
  289. print("Attack number: " + str(i))
  290. print("Attack name: " + str(label.attack_name))
  291. print("Attack note: " + str(label.attack_note))
  292. print("Attack seed: " + str(label.seed))
  293. print("Start timestamp: " + str(label.timestamp_start))
  294. print("End timestamp: " + str(label.timestamp_end))
  295. print()
  296. print()
  297. return True
  298. elif cmd == "set":
  299. if len(params) == 3:
  300. if params[0].lower() == "attack_note":
  301. i = int(params[1])
  302. self.label_manager.labels[i].attack_note = params[2]
  303. return True
  304. elif cmd == "tables":
  305. self.statisticsDB.process_db_query("SELECT name FROM sqlite_master WHERE type='table';", True)
  306. return True
  307. elif cmd == "columns":
  308. is_table = self.statisticsDB.process_db_query("SELECT name FROM sqlite_master WHERE type='table' AND name='"
  309. + params[0].lower() + "';", False)
  310. if not is_table:
  311. print("Table " + params[0].lower() + " does not exist.")
  312. return True
  313. self.statisticsDB.process_db_query("SELECT * FROM " + params[0].lower(), False)
  314. columns = self.statisticsDB.get_field_types(params[0].lower())
  315. for column in columns:
  316. print(column + ": " + columns[column])
  317. return True
  318. return False
  319. def enter_query_mode(self):
  320. """
  321. Enters into the query mode. This is a read-eval-print-loop, where the user can input named queries or SQL
  322. queries and the results are printed.
  323. """
  324. def make_completer(vocabulary):
  325. """
  326. TODO: FILL ME
  327. :param vocabulary:
  328. :return:
  329. """
  330. def custom_template(text, state):
  331. """
  332. TODO: FILL ME
  333. :param text:
  334. :param state:
  335. :return:
  336. """
  337. results = [x for x in vocabulary if x.startswith(text)] + [None]
  338. return results[state]
  339. return custom_template
  340. readline.parse_and_bind('tab: complete')
  341. readline.set_completer(make_completer(
  342. self.statisticsDB.get_all_named_query_keywords() + self.statisticsDB.get_all_sql_query_keywords()))
  343. history_file = os.path.join(Util.CACHE_DIR, 'query_history')
  344. try:
  345. readline.read_history_file(history_file)
  346. except IOError:
  347. pass
  348. print("Entering into query mode...")
  349. print("Enter statement ending by ';' and press ENTER to send query. Exit by sending an empty query.")
  350. print("Type 'help;' for information on possible queries.")
  351. buffer = ""
  352. while True:
  353. line = input("> ")
  354. if line == "":
  355. break
  356. buffer += line
  357. import sqlite3
  358. if sqlite3.complete_statement(buffer):
  359. buffer = buffer.strip()
  360. if not self.internal_command(buffer):
  361. try:
  362. self.statisticsDB.process_db_query(buffer, True)
  363. except sqlite3.Error as e:
  364. print("An error occurred:", e.args[0])
  365. except pp.ParseException as e:
  366. sys.stderr.write("Error in query:\n")
  367. sys.stderr.write(buffer)
  368. sys.stderr.write("\n")
  369. for i in range(1, e.col):
  370. sys.stderr.write(" ")
  371. sys.stderr.write("^\n\n")
  372. except StatsDB.QueryExecutionException as e:
  373. sys.stderr.write("An error occured: ")
  374. sys.stderr.write(e.args[0] + "\n")
  375. buffer = ""
  376. readline.set_history_length(1000)
  377. readline.write_history_file(history_file)
  378. # Save the label file, in case content has changed
  379. self.label_manager.write_label_file(self.pcap_src_path)
  380. def create_statistics_plot(self, params: str, entropy: bool):
  381. """
  382. Plots the statistics to a file by using the given customization parameters.
  383. """
  384. print("Statistical plots are being generated", end="", flush=True)
  385. if params is not None and params[0] is not None:
  386. # FIXME: cleanup
  387. params_dict = dict([z.split("=") for z in params])
  388. self.statistics.plot_statistics(entropy=entropy, file_format=params_dict['format'])
  389. else:
  390. self.statistics.plot_statistics(entropy=entropy)