def main_command_line(): # Since plugin_chain contains the actual plugin instances we have to make sure # we reset the global plugin_chain so multiple runs don't affect each other. # (This was necessary to call this function through a python script.) # TODO: Should plugin_chain be a list of plugin classes instead of instances? global plugin_chain plugin_chain = [] # dictionary of all available plugins: {name: module path} plugin_map = get_plugins() # dictionary of plugins that the user wants to use: {name: object} active_plugins = OrderedDict() # The main argument parser. It will have every command line option # available and should be used when actually parsing parser = DshellArgumentParser( usage="%(prog)s [options] [plugin options] file1 file2 ... fileN", add_help=False) parser.add_argument('-c', '--count', type=int, default=0, help='Number of packets to process') parser.add_argument('--debug', action="store_true", help="Show debug messages") parser.add_argument('-v', '--verbose', action="store_true", help="Show informational messages") parser.add_argument( '-acc', '--allcc', action="store_true", help= "Show all 3 GeoIP2 country code types (represented_country/registered_country/country)" ) parser.add_argument( '-d', '-p', '--plugin', dest='plugin', type=str, action='append', metavar="PLUGIN", help="Use a specific plugin module. Can be chained with '+'.") parser.add_argument('--defragment', dest='defrag', action='store_true', help='Reconnect fragmented IP packets') parser.add_argument('-h', '-?', '--help', dest='help', help="Print common command-line flags and exit", action='store_true', default=False) parser.add_argument( '-i', '--interface', default=None, type=str, help="Listen live on INTERFACE instead of reading pcap") parser.add_argument('-l', '--ls', '--list', action="store_true", help='List all available plugins', dest='list') parser.add_argument( '-r', '--recursive', dest='recursive', action='store_true', help='Recursively process all PCAP files under input directory') parser.add_argument( '--unzipdir', type=str, metavar="DIRECTORY", default=tempfile.gettempdir(), help= 'Directory to use when decompressing input files (.gz, .bz2, and .zip only)' ) multiprocess_group = parser.add_argument_group("multiprocessing arguments") multiprocess_group.add_argument( '-P', '--parallel', dest='multiprocessing', action='store_true', help='Handle each file in separate parallel processes') multiprocess_group.add_argument( '-n', '--nprocs', type=int, default=4, metavar='NUMPROCS', dest='process_max', help='Define max number of parallel processes (default: 4)') filter_group = parser.add_argument_group("filter arguments") filter_group.add_argument( '--bpf', default='', type=str, help="Overwrite all BPFs and use provided input. Use carefully!") filter_group.add_argument( '--ebpf', default='', type=str, metavar="BPF", help= "Extend existing BPFs with provided input for additional filtering. It will transform input into \"(<original bpf>) and (<ebpf>)\"" ) filter_group.add_argument("--no-vlan", action="store_true", dest="novlan", help="Ignore packets with VLAN headers") output_group = parser.add_argument_group("output arguments") output_group.add_argument("--lo", "--list-output", action="store_true", help="List available output modules", dest="listoutput") output_group.add_argument("--no-buffer", action="store_true", help="Do not buffer plugin output", dest="nobuffer") output_group.add_argument("-x", "--extra", action="store_true", help="Appends extra data to all plugin output.") # TODO Figure out how to make --extra flag play nicely with user-only # output modules, like jsonout and csvout output_group.add_argument( "-O", "--omodule", type=str, dest="omodule", metavar="MODULE", help= "Use specified output module for plugins instead of defaults. For example, --omodule=jsonout for JSON output." ) output_group.add_argument( "--oarg", type=str, metavar="ARG=VALUE", dest="oargs", action="append", help= "Supply a specific keyword argument to plugins' output modules. Can be used multiple times for multiple arguments. Not using an equal sign will treat it as a flag and set the value to True. Example: --oarg \"delimiter=:\" --oarg \"timeformat=%%H %%M %%S\"" ) output_group.add_argument("-q", "--quiet", action="store_true", help="Disable logging") output_group.add_argument("-W", metavar="OUTFILE", dest="outfile", help="Write to OUTFILE instead of stdout") parser.add_argument('files', nargs='*', help="pcap files or globs to process") # A short argument parser, meant to only hold the simplified list of # arguments for when a plugin is called without a pcap file. # DO NOT USE for any serious argument parsing. parser_short = DshellArgumentParser( usage="%(prog)s [options] [plugin options] file1 file2 ... fileN", add_help=False) parser_short.add_argument('-h', '-?', '--help', dest='help', help="Print common command-line flags and exit", action='store_true', default=False) parser.add_argument('--version', action='version', version="Dshell " + str(dshell.core.__version__)) parser_short.add_argument('-d', '-p', '--plugin', dest='plugin', type=str, action='append', metavar="PLUGIN", help="Use a specific plugin module") parser_short.add_argument( '--ebpf', default='', type=str, metavar="BPF", help= "Extend existing BPFs with provided input for additional filtering. It will transform input into \"(<original bpf>) and (<ebpf>)\"" ) parser_short.add_argument( '-i', '--interface', help="Listen live on INTERFACE instead of reading pcap") parser_short.add_argument('-l', '--ls', '--list', action="store_true", help='List all available plugins', dest='list') parser_short.add_argument("--lo", "--list-output", action="store_true", help="List available output modules") # FIXME: Should this duplicate option be removed? parser_short.add_argument( "-o", "--omodule", type=str, metavar="MODULE", help= "Use specified output module for plugins instead of defaults. For example, --omodule=jsonout for JSON output." ) parser_short.add_argument('files', nargs='*', help="pcap files or globs to process") # Start parsing the arguments # Specifically, we want to grab the desired plugin list # This will let us add the plugin-specific arguments and reprocess the args opts, xopts = parser.parse_known_args() if opts.plugin: # Multiple plugins can be chained using either multiple instances # of -d/-p/--plugin or joining them together with + signs. plugins = '+'.join(opts.plugin) plugins = plugins.split('+') # check for invalid plugins for plugin in plugins: plugin = plugin.strip() if not plugin: # User probably mistyped '++' instead of '+' somewhere. # Be nice and ignore this minor infraction. continue if plugin not in plugin_map: parser_short.epilog = "ERROR! Invalid plugin provided: '{}'".format( plugin) parser_short.print_help() sys.exit(1) # While we're at it, go ahead and import the plugin modules now # This can probably be done further down the line, but here is # just convenient plugin_module = import_module(plugin_map[plugin]) # Handle multiple instances of same plugin by appending number to # end of plugin name. This is used mostly to separate # plugin-specific arguments from each other if plugin in active_plugins: i = 1 plugin = plugin + str(i) while plugin in active_plugins: i += 1 plugin = plugin[:-(len(str(i - 1)))] + str(i) # Add copy of plugin object to chain and add to argument parsers # TODO: Use class attributes for class related things like name, description, optionsdict # This way we don't have to initialize the plugin at this point and fixes a lot of the # issues that arise that come from dealing with a singleton. active_plugins[plugin] = plugin_module.DshellPlugin() plugin_chain.append(active_plugins[plugin]) parser.add_plugin_arguments(plugin, active_plugins[plugin]) parser_short.add_plugin_arguments(plugin, active_plugins[plugin]) opts, xopts = parser.parse_known_args() if xopts: for xopt in xopts: logger.warning('Could not understand argument {!r}'.format(xopt)) if opts.help: # Just print the full help message and exit parser.print_help() print("\n") for plugin in plugin_chain: print("############### {}".format(plugin.name)) print(plugin.longdescription) print("\n") print('Default BPF: "{}"'.format(plugin.bpf)) print("\n") sys.exit() if opts.list: try: print_plugins(get_plugin_information()) except ImportError as e: logger.error(e, exc_info=opts.debug) sys.exit() if opts.listoutput: # List available output modules and a brief description output_map = get_output_modules(get_output_path()) for modulename in sorted(output_map): try: module = import_module("dshell.output." + modulename) module = module.obj except Exception as e: etype = e.__class__.__name__ logger.debug("Could not load {} module. ({}: {!s})".format( modulename, etype, e)) else: print("\t{:<25} {}".format(modulename, module._DESCRIPTION)) sys.exit() if not opts.plugin: # If a plugin isn't provided, print the short help message parser_short.epilog = "Select a plugin to use with -d or --plugin" parser_short.print_help() sys.exit() if not opts.files and not opts.interface: # If no files are provided, print the short help message parser_short.epilog = "Include a pcap file to get started. Use --help for more information." parser_short.print_help() sys.exit() # Process the plugin-specific args and set the attributes within them plugin_args = {} for plugin_name, plugin in active_plugins.items(): plugin_args[plugin] = {} args_and_attrs = parser.get_plugin_arguments(plugin_name, plugin) for darg, dattr in args_and_attrs: value = getattr(opts, darg) plugin_args[plugin][dattr] = value main(plugin_args=plugin_args, **vars(opts))
def main(plugin_args=None, **kwargs): global plugin_chain if not plugin_args: plugin_args = {} # dictionary of all available plugins: {name: module path} plugin_map = get_plugins() # Attempt to catch segfaults caused when certain linktypes (e.g. 204) are # given to pcapy faulthandler.enable() if not plugin_chain: logger.error("No plugin selected") sys.exit(1) plugin_chain[0].defrag_ip = kwargs.get("defrag", False) # Setup logging log_format = "%(levelname)s (%(name)s) - %(message)s" if kwargs.get("verbose", False): log_level = logging.INFO elif kwargs.get("debug", False): log_level = logging.DEBUG elif kwargs.get("quiet", False): log_level = logging.CRITICAL else: log_level = logging.WARNING logging.basicConfig(format=log_format, level=log_level) # since pypacker handles its own exceptions (loudly), this attempts to keep # it quiet logging.getLogger("pypacker").setLevel(logging.CRITICAL) if kwargs.get("allcc", False): # Activate all country code (allcc) mode to display all 3 GeoIP2 country # codes dshell.core.geoip.acc = True dshell.core.geoip.check_file_dates() # If alternate output module is selected, tell each plugin to use that # instead if kwargs.get("omodule", None): try: # TODO: Create a factory classmethod in the base Output class (e.g. "from_name()") instead. omodule = import_module("dshell.output." + kwargs["omodule"]) omodule = omodule.obj for plugin in plugin_chain: # TODO: Should we have a single instance of the Output module used by all plugins? oomodule = omodule() plugin.out = oomodule except ImportError as e: logger.error( "Could not import module named '{}'. Use --list-output flag to see available modules" .format(kwargs["omodule"])) sys.exit(1) # Check if any user-defined output arguments are provided if kwargs.get("oargs", None): oargs = {} for oarg in kwargs["oargs"]: if '=' in oarg: key, val = oarg.split('=', 1) oargs[key] = val else: oargs[oarg] = True logger.debug("oargs: %s" % oargs) for plugin in plugin_chain: plugin.out.set_oargs(**oargs) # If writing to a file, set for each output module here if kwargs.get("outfile", None): for plugin in plugin_chain: plugin.out.reset_fh(filename=kwargs["outfile"]) # Set nobuffer mode if that's what the user wants if kwargs.get("nobuffer", False): for plugin in plugin_chain: plugin.out.nobuffer = True # Set the extra flag for all output modules if kwargs.get("extra", False): for plugin in plugin_chain: plugin.out.extra = True plugin.out.set_format(plugin.out.format) # Set the BPF filters # Each plugin has its own default BPF that will be extended or replaced # based on --no-vlan, --ebpf, or --bpf arguments. for plugin in plugin_chain: if kwargs.get("bpf", None): plugin.bpf = kwargs.get("bpf", "") continue if plugin.bpf: if kwargs.get("ebpf", None): plugin.bpf = "({}) and ({})".format(plugin.bpf, kwargs.get("ebpf", "")) else: if kwargs.get("ebpf", None): plugin.bpf = kwargs.get("ebpf", "") if kwargs.get("novlan", False): plugin.vlan_bpf = False # Decide on the inputs to use for pcap # If --interface is set, ignore all files and listen live on the wire # Otherwise, use all of the files and globs to open offline pcap. # Recurse through any directories if the command-line flag is set. if kwargs.get("interface", None): inputs = [kwargs.get("interface")] else: inputs = [] inglobs = kwargs.get("files", []) infiles = [] for inglob in inglobs: outglob = glob(inglob) if not outglob: logger.warning( "Could not find file(s) matching {!r}".format(inglob)) continue infiles.extend(outglob) while len(infiles) > 0: infile = infiles.pop(0) if kwargs.get("recursive", False) and os.path.isdir(infile): morefiles = os.listdir(infile) for morefile in morefiles: infiles.append(os.path.join(infile, morefile)) elif os.path.isfile(infile): inputs.append(infile) # Process plugin-specific options for plugin in plugin_chain: for option, args in plugin.optiondict.items(): if option in plugin_args.get(plugin, {}): setattr(plugin, option, plugin_args[plugin][option]) else: setattr(plugin, option, args.get("default", None)) plugin.handle_plugin_options() #### Dshell is ready to read pcap! #### for plugin in plugin_chain: plugin._premodule() # If we are not multiprocessing, simply pass the files for processing if not kwargs.get("multiprocessing", False): process_files(inputs, **kwargs) # If we are multiprocessing, things get more complicated. else: # Create an output queue, and wrap the 'write' function of each # plugins's output module to send calls to the multiprocessing queue output_queue = multiprocessing.Queue() output_wrappers = {} for plugin in plugin_chain: qo = QueueOutputWrapper(plugin.out, output_queue) output_wrappers[qo.id] = qo plugin.out.write = qo.write # Create processes to handle each separate input file processes = [] for i in inputs: processes.append( multiprocessing.Process(target=process_files, args=([i], ), kwargs=kwargs)) # Spawn processes, and keep track of which ones are running running = [] max_writes_per_batch = 50 while processes or running: if processes and len(running) < kwargs.get("process_max", 4): # Start a process and move it to the 'running' list proc = processes.pop(0) proc.start() logger.debug("Started process {}".format(proc.pid)) running.append(proc) for proc in running: if not proc.is_alive(): # Remove finished processes from 'running' list logger.debug("Ended process {} (exit code: {})".format( proc.pid, proc.exitcode)) running.remove(proc) try: # Process write commands in the output queue. # Since some plugins write copiously and may block other # processes from launching, only write up to a maximum number # before breaking and rechecking the processes. writes = 0 while writes < max_writes_per_batch: wrapper_id, args, kwargs = output_queue.get(True, 1) owrapper = output_wrappers[wrapper_id] owrapper.true_write(*args, **kwargs) writes += 1 except queue.Empty: pass output_queue.close() for plugin in plugin_chain: plugin._postmodule()