CLI.py 8.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194
  1. #! /usr/bin/env python3
  2. import argparse
  3. import sys
  4. import random
  5. import numpy
  6. import hashlib
  7. from ID2TLib.Controller import Controller
  8. class LoadFromFile(argparse.Action):
  9. """
  10. Parses the parameter file given by application param -c/--config.
  11. """
  12. def __call__(self, parser, namespace, values, option_string=None):
  13. with values as f:
  14. parser.parse_args(f.read().split(), namespace)
  15. class CLI(object):
  16. def __init__(self):
  17. """
  18. Creates a new CLI object used to handle
  19. """
  20. # Reference to PcapFile object
  21. self.args = None
  22. self.attack_config = None
  23. def parse_arguments(self, args):
  24. """
  25. Defines the allowed application arguments and invokes the evaluation of the arguments.
  26. :param args: The application arguments
  27. """
  28. # Create parser for arguments
  29. parser = argparse.ArgumentParser(description="Intrusion Detection Dataset Toolkit (ID2T) - A toolkit for "
  30. "injecting synthetically created attacks into PCAP files.",
  31. prog="id2t")
  32. # Required arguments
  33. required_group = parser.add_argument_group('required arguments')
  34. required_args_group = required_group.add_mutually_exclusive_group(required=True)
  35. required_args_group.add_argument('-i', '--input', metavar="PCAP_FILE",
  36. help='path to the input pcap file')
  37. required_args_group.add_argument('-l', '--list-attacks', action='store_true')
  38. # Optional arguments
  39. parser.add_argument('-c', '--config', metavar='CONFIG_FILE', help='file containing configuration parameters.',
  40. action=LoadFromFile, type=open)
  41. parser.add_argument('-e', '--export',
  42. help='store statistics as a ".stat" file',
  43. action='store_true', default=False)
  44. parser.add_argument('-r', '--recalculate',
  45. help='recalculate statistics even if a cached version exists.',
  46. action='store_true', default=False)
  47. parser.add_argument('-s', '--statistics', help='print file statistics to stdout.', action='store_true',
  48. default=False)
  49. parser.add_argument('-p', '--plot', help='creates the following plots: the values distributions of TTL, MSS, Window Size, '
  50. 'protocol, and the novelty distributions of IP, port, TTL, MSS, Window Size,'
  51. ' and ToS. In addition to packets count in interval-wise.', action='append',
  52. nargs='?')
  53. parser.add_argument('-q', '--query', metavar="QUERY",
  54. action='append', nargs='?',
  55. help='query the statistics database. If no query is provided, the application enters query mode.')
  56. parser.add_argument('-t', '--extraTests', help='perform extra tests on the input pcap file, including calculating IP entropy'
  57. 'in interval-wise, TCP checksum, and checking payload availability.', action='store_true')
  58. parser.add_argument('--seed', help='random seed for the program, call the program with the same seed to get the same output')
  59. parser.add_argument('-o', '--output', metavar="PCAP_FILE",
  60. help='path to the output pcap file')
  61. parser.add_argument('-ie', '--inject_empty', action='store_true',
  62. help='injects ATTACK into an EMPTY PCAP file, using the statistics of the input PCAP.')
  63. # Attack arguments
  64. parser.add_argument('-a', '--attack', metavar="ATTACK", action='append',
  65. help='injects ATTACK into a PCAP file.', nargs='+')
  66. # Parse arguments
  67. self.args = parser.parse_args(args)
  68. self.process_arguments()
  69. def process_arguments(self):
  70. """
  71. Decide what to do with each of the command line parameters.
  72. """
  73. if self.args.seed is not None:
  74. self.seed_rng(self.args.seed)
  75. if self.args.list_attacks:
  76. # User wants to see the available attacks
  77. self.process_attack_listing()
  78. else:
  79. # User wants to process a PCAP
  80. self.process_pcap()
  81. def seed_rng(self, seed):
  82. try: # try to convert the seed to int
  83. seed = int(seed)
  84. except: # otherwise use the strings hash
  85. hashed_seed = hashlib.sha1(seed.encode()).digest()
  86. seed = int.from_bytes(hashed_seed, byteorder="little") & 0xffffffff # convert hash to 32-bit integer
  87. random.seed(seed)
  88. numpy.random.seed(seed)
  89. def process_attack_listing(self):
  90. import pkgutil
  91. import importlib
  92. import Attack
  93. # Find all attacks, exclude some classes
  94. package = Attack
  95. attack_names = []
  96. for _, name, __ in pkgutil.iter_modules(package.__path__):
  97. if name != 'BaseAttack' and name != 'AttackParameters':
  98. attack_names.append(name)
  99. # List the attacks and their parameters
  100. emph_start = '\033[1m'
  101. emph_end = '\033[0m'
  102. for attack_name in attack_names:
  103. attack_module = importlib.import_module('Attack.{}'.format(attack_name))
  104. attack_class = getattr(attack_module, attack_name)
  105. # Instantiate the attack to get to its definitions.
  106. attack_obj = attack_class()
  107. print('* {}{}{}'.format(emph_start, attack_obj.attack_name, emph_end))
  108. print('\t- {}Description:{} {}'.format(emph_start, emph_end,
  109. attack_obj.attack_description))
  110. print('\t- {}Type:{} {}'.format(emph_start, emph_end,
  111. attack_obj.attack_type))
  112. print('\t- {}Supported Parameters:{}'.format(emph_start, emph_end), end=' ')
  113. # Get all the parameter names in a list and sort them
  114. param_list = []
  115. for key in attack_obj.supported_params:
  116. param_list.append(key.value)
  117. param_list.sort()
  118. # Print each parameter type per line
  119. last_prefix = None
  120. current_prefix = None
  121. for param in param_list:
  122. current_prefix = param.split('.')[0]
  123. if not last_prefix or current_prefix != last_prefix:
  124. print('\n\t + |', end=' ')
  125. print(param, end=' | ')
  126. last_prefix = current_prefix
  127. # Print an empty line
  128. print()
  129. def process_pcap(self):
  130. """
  131. Loads the application controller, the PCAP file statistics and if present, processes the given attacks. Evaluates
  132. given queries.
  133. """
  134. # Create ID2T Controller
  135. controller = Controller(self.args.input, self.args.extraTests, self.args.output)
  136. # Load PCAP statistics
  137. controller.load_pcap_statistics(self.args.export, self.args.recalculate, self.args.statistics)
  138. # Create statistics plots
  139. if self.args.plot is not None:
  140. controller.create_statistics_plot(self.args.plot)
  141. # Process attack(s) with given attack params
  142. if self.args.attack is not None:
  143. # If attack is present, load attack with params
  144. if self.args.inject_empty:
  145. # Attack PCAPs will not be merged with base PCAP
  146. controller.process_attacks(self.args.attack, inject_empty=True)
  147. else:
  148. # Attack PCAP will be merged with base PCAP
  149. controller.process_attacks(self.args.attack)
  150. # Parameter -q without arguments was given -> go into query loop
  151. if self.args.query == [None]:
  152. controller.enter_query_mode()
  153. # Parameter -q with arguments was given -> process query
  154. elif self.args.query is not None:
  155. controller.process_db_queries(self.args.query, True)
  156. def main(args):
  157. """
  158. Creates a new CLI object and invokes the arguments parsing.
  159. :param args: The provided arguments
  160. """
  161. cli = CLI()
  162. # Check arguments
  163. cli.parse_arguments(args)
  164. # Uncomment to enable calling by terminal
  165. if __name__ == '__main__':
  166. main(sys.argv[1:])