Example #1
0
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))
Example #2
0
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()