All VxStream related stuff may be removed in a future version

def main():

    parser = argparse.ArgumentParser(description="An interface to CarbonBlack environments")

    #profiles = auth.CredentialStore("response").get_profiles()
    parser.add_argument('-e', '--environment', choices=auth.CredentialStore("response").get_profiles(),
                        help='specify a specific instance you want to work with. If not defined \'-t production\' will be used implicitly.')
    parser.add_argument('-t', '--envtypes', type=str, 
                        help='specify any combination of envtypes. Default=All \'production\' envtypes. Ignored if -e is set.',
                        default='production')
    #parser.add_argument('--debug', action='store_true', help='print debugging info')
    #parser.add_argument('--warnings', action='store_true',
    #                         help="Warn before printing large executions")

    subparsers = parser.add_subparsers(dest='command') #title='subcommands', help='additional help')
    cbinterface_commands = [ 'query', 'proc', 'collect', 'remediate', 'enumerate_usb', 'vxdetect']

    parser_vx = subparsers.add_parser('vxdetect', help="search cbsandbox for processes in vxstream report, show detections")
    parser_vx.add_argument('vxstream_report', help='path to vxstream report')
    parser_vx.add_argument('-p', '--print-process-tree', action='store_true', help='print the process tree')

    parser_usb = subparsers.add_parser('enumerate_usb', help="Show recent removable drive activity on the sensor")
    parser_usb.add_argument('sensor', help='hostname of the sensor')
    parser_usb.add_argument('-s', '--start-time', action='store',
                            help='how far back to query (default:ALL time)')

    parser_proc = subparsers.add_parser('proc', help="analyze a process GUID. 'proc -h' for more")
    parser_proc.add_argument('process', help="the process GUID to analyze")
    parser_proc.add_argument('--warnings', action='store_true',
                             help="Warn before printing large executions")
    parser_proc.add_argument('-w', '--walk-tree', action='store_true',
                             help="walk and analyze the process tree")
    parser_proc.add_argument('-wp', '--walk-parents', action='store_true',
                             help="print details on the process ancestry")
    #parser_proc.add_argument('-d', '--detection', action='store_true',
    #                         help="show detections that would result in ACE alerts")
    parser_proc.add_argument('-i', '--proc-info', action='store_true',
                             help="show binary and process information")
    parser_proc.add_argument('-c','--show-children', action='store_true',
                             help="only print process children event details")
    parser_proc.add_argument('-nc', '--netconns', action='store_true',
                             help="print network connections")
    parser_proc.add_argument('-fm', '--filemods', action='store_true',
                             help="print file modifications")
    parser_proc.add_argument('-rm', '--regmods', action='store_true',
                             help="print registry modifications")
    parser_proc.add_argument('-um', '--unsigned-modloads', action='store_true',
                             help="print unsigned modloads")
    parser_proc.add_argument('-ml', '--modloads', action='store_true',
                             help="print modloads")
    parser_proc.add_argument('-cp', '--crossprocs', action='store_true',
                             help="print crossprocs")
    #parser_proc.add_argument('-intel', '--intel-hits', action='store_true',
    #                         help="show intel (feed/WL) hits that do not result in ACE alerts")
    parser_proc.add_argument('--no-analysis', action='store_true',
                             help="Don't fetch and print process activity")
    parser_proc.add_argument('--json', action='store_true', help='output process summary in json')
    parser_proc.add_argument('--segment-limit', action='store', type=int, default=None,
                             help='stop processing events into json after this many process segments')

    facet_args = [
        'process_name', 'childproc_name', 'username', 'parent_name', 'path', 'hostname',
        'parent_pid', 'comms_ip', 'process_md5', 'start', 'group', 'interface_ip',
        'modload_count', 'childproc_count', 'cmdline', 'regmod_count', 'process_pid',
        'parent_id', 'os_type', 'rocessblock_count', 'crossproc_count', 'netconn_count',
        'parent_md5', 'host_type', 'last_update', 'filemod_count'
        ]
 
    parser_query = subparsers.add_parser('query',
                                         help="execute a process search query. 'query -h' for more")
    parser_query.add_argument('query', help="the process search query you'd like to execute")
    parser_query.add_argument('-s', '--start-time', action='store',
                              help="Only return processes with events after given date/time stamp\
 (server’s clock). Format:'Y-m-d H:M:S' eastern time")
    parser_query.add_argument('-e', '--end-time', action='store',
                              help="Set the maximum last update time. Format:'Y-m-d H:M:S' eastern time")
    parser_query.add_argument('--facet', action='store', choices=facet_args,
                              help='stats info on single field accross query results (ex. process_name)')
    parser_query.add_argument('--no-warnings', action='store_true',
                             help="Don't warn before printing large query results")
    parser_query.add_argument('-lh', '--logon-history', action='store_true', help="Display available logon history for given username or hostname")

    parser_collect = subparsers.add_parser('collect', help='perform LR collection tasks on a host')
    parser_collect.add_argument('sensor', help="the hostname/sensor to collect from")
    parser_collect.add_argument('-f', '--filepath', action='store', help='collect file')
    parser_collect.add_argument('-c', '--command-exec', action='store', help='command to execute')
    parser_collect.add_argument('-p', '--process-list', action='store_true', 
                                help='show processes running on sensor')
    parser_collect.add_argument('-m', '--memdump', action='store', const='ALLMEM', nargs='?',
                                help='dump memory on a specific process-id')
    parser_collect.add_argument('-lr', '--regkeypath', action='store',
                                help='List all registry values from the specified registry key.')
    parser_collect.add_argument('-r', '--regkeyvalue', action='store',
                                help='Returns the associated value of the specified registry key.')
    parser_collect.add_argument('-i', '--info', action='store_true', help='print sensor information')
    parser_collect.add_argument('-gst', '--get-task', action='store_true', help='get scheduled tasks or specifc task')
    parser_collect.add_argument('-mc', '--multi-collect', action='store', help='path to ini file listing files and regs to collect')

    remediate_file_example = """Example remediate ini file:
    [files]
    file1=C:\\Users\\user\\Desktop\\testfile.txt 

    [process_names]
    proc1=cmd.exe
    proc2=notepad++.exe
     
    [directories]
    directory1=C:\\Users\\user\\Desktop\\nanocore
     
    [scheduled_tasks]
    task1=\\monkey_task
    task1=\\Microsoft\\windows\\some\\flake\\task

    [pids]
    pid1=10856
     
    [registry_paths]
    reg1=HKLM\\Software\\Microsoft\\Windows\\CurrentVersion\\Run\\calc
    reg2=HKLM\\Software\\Microsoft\\Windows\\CurrentVersion\\Run\\hippo"""

    parser_remediate = subparsers.add_parser('remediate', help='remediate a host')
    parser_remediate.add_argument('sensor', help="the hostname/sensor needing remediation")
    parser_remediate.add_argument('-i', '--isolate', help='toggle host isolation', default=False, action='store_true')
    parser_remediate.add_argument('-f', '--remediation-filepath',
                                  help="path to the remediation ini file; 'help' as the filepath for example")
    parser_remediate.add_argument('-dst', '--delete-scheduled-task',
                                  help="path of scheduled task to delete")
    parser_remediate.add_argument('-kpname', '--kill-process-name', help="kill all processes with this name")
    parser_remediate.add_argument('-kpid', '--kill-pid', help="a process id to kill")
    parser_remediate.add_argument('-df', '--delete-file', help="path to file needing deletion")
    parser_remediate.add_argument('-dr', '--delete-regkey', help="path to regkey value needing deletion")
    parser_remediate.add_argument('-dd', '--delete-directory', help="path to directory needing deletion")
    args = parser.parse_args()

    if args.command == 'remediate' and args.remediation_filepath == 'help':
        print(remediate_file_example)
        parser.parse_args(['remediate', '-h'])

    if args.command is None:
        print("\n\n*****")
        print("You must specify one of the following commands:\n")
        print(cbinterface_commands)
        print("\n*****\n\n")
        parser.parse_args(['-h'])

    #args.debug = True
    #if args.debug:
    # configure some more logging
    root = logging.getLogger()
    root.addHandler(logging.StreamHandler())
    logging.getLogger("cbapi").setLevel(logging.ERROR)
    logging.getLogger("lerc_api").setLevel(logging.WARNING)

    ''' All VxStream related stuff may be removed in a future version '''
    if args.command == 'vxdetect':
        cb = CbResponseAPI(profile='vxstream')
        process_list = parse_vxstream_report(cb, args.vxstream_report)
        if args.print_process_tree:
            print()
            print(process_list)
        print()
        return 0


    # Set up environment profiles
    profile = None
    profiles = []
    if args.environment: 
        print("Using {} environment ..".format(args.environment))
        profiles.append(args.environment)
    else:
        # a little hack for getting our environment type variable defined
        default_profile = auth.default_profile
        default_profile['envtype'] = 'production'
        query_envtype = set(args.envtypes.lower().split(','))
        for profile in auth.CredentialStore("response").get_profiles():
            credentials = auth.CredentialStore("response").get_credentials(profile=profile)
            profile_envtype = set(credentials['envtype'].lower().split(','))
            if(query_envtype.issubset(profile_envtype)):
                profiles.append(profile)
        
        


    # Process Quering #
    if args.command == 'query':
        for profile in profiles:
            handle_proxy(profile)
            print("\nSearching {} environment..".format(profile))
            q = CBquery(profile=profile)
            q.process_query(args)
        return 0


    # Select correct environment by sensor hostname and get the sensor object
    sensor = None
    if args.command == 'collect' or args.command == 'remediate' or args.command == 'enumerate_usb':
        cb_results = sensor_search(profiles, args.sensor)
        if not isinstance(cb_results, list):
            # an error occured
            return cb_results
        else:
            if not cb_results:
                LOGGER.info("A sensor with hostname {} wasn't found in any environments".format(args.sensor))
                return 0
            elif len(cb_results) > 1:
                LOGGER.error("A sensor by hostname {} was found in multiple environments".format(args.sensor))
                for r in cb_results:
                    print("Results:")
                    print("Profile {}: {} (SID:{})".format(r[1],r[0].hostname,r[0].id))
                return 1
            results = cb_results[0]
            profile = results[1]
            sensor = results[0]


    # Show USB Regmod activity
    if args.command == 'enumerate_usb':
        enumerate_usb(sensor, args.start_time)


    # lerc install arguments can differ by company/environment
    # same lazy hack to define in cb config
    config = {}
    try:
        default_profile = auth.default_profile
        default_profile['lerc_install_cmd'] = None
        config = auth.CredentialStore("response").get_credentials(profile=profile)
    except:
        pass


    # Collection #
    if args.command == 'collect':
        hyper_lr = hyperLiveResponse(sensor)

        if args.info:
            print(hyper_lr)
            return True

        # start a cb lr session
        lr_session = hyper_lr.go_live()

        if args.multi_collect:
            filepaths = regpaths = full_collect = None
            config = ConfigParser()
            config.read(args.multi_collect)
            try:
                filepaths = config.items("files")
            except:
                filepaths = []
            try:
                regpaths = config.items("registry_paths")
            except:
                regpaths = []
            try:
                full_collect = config.get('full_collect', 'action')
            except:
                pass

            if regpaths is not None:
                for regpath in regpaths:
                    if isinstance(regpath, tuple):
                        regpath = regpath[1]
                    print("~ Trying to get {}".format(regpath))
                    try:
                        result = lr_session.get_registry_value(regpath)
                        if result:
                            localfname = args.sensor + '_regkey_' + result['value_name'] + ".txt"
                            with open(localfname,'wb') as f:
                                f.write(bytes(result['value_data'], 'UTF-8'))
                            print("\t+ Data written to: {}".format(localfname))
                    except Exception as e:
                        print("[!] Error: {}".format(str(e)))
            if filepaths is not None:
                for filepath in filepaths:
                    try:
                        hyper_lr.getFile_with_timeout(filepath[1])
                    except Exception as e:
                        print("[!] Error: {}".format(str(e)))
            if full_collect == 'True':
               return False #LR_collection(hyper_lr, args)
            return True

        elif args.filepath:
            hyper_lr.getFile_with_timeout(args.filepath)

        elif args.process_list:
            hyper_lr.print_processes()

        elif args.memdump:
            # get config items
            config = ConfigParser()
            config.read(CONFIG_PATH)
            #if config.has_section('memory'):
            #    if 
            cb_compress = config['memory'].getboolean('cb_default_compress')
            custom_compress = config['memory'].getboolean('custom_compress')
            custom_compress_file = config['memory']['custom_compress_file']
            auto_collect_mem = config['memory'].getboolean('auto_collect_mem_file')
            lerc_collect_mem = config['memory'].getboolean('lerc_collect_mem')
            path_to_procdump = config['memory']['path_to_procdump']
            
            if args.memdump == "ALLMEM":
                return hyper_lr.dump_sensor_memory(cb_compress=cb_compress, custom_compress=custom_compress,
                                                   custom_compress_file=custom_compress_file,
                                                   auto_collect_result=auto_collect_mem)
            else:
                return hyper_lr.dump_process_memory(args.memdump, path_to_procdump=path_to_procdump)

        elif args.command_exec:
            print("executing '{}' on {}".format(args.command_exec, args.sensor))
            result = lr_session.create_process(args.command_exec, wait_timeout=60, wait_for_output=True)
            print("\n-------------------------")
            result = result.decode('utf-8')
            print(result + "\n-------------------------")
            print()

        elif args.regkeypath:
            print("\n\t{}".format(args.regkeypath))
            results = lr_session.list_registry_keys(args.regkeypath)
            for result in results:
                print("\t-------------------------")
                print("\tName: {}".format(result['value_name']))
                print("\tType: {}".format(result['value_type']))
                print("\tData: {}".format(result['value_data']))
                print()

        elif args.regkeyvalue:
            print("\n\t{}".format(args.regkeyvalue))
            result = lr_session.get_registry_value(args.regkeyvalue)
            print("\t-------------------------")
            print("\tName: {}".format(result['value_name']))
            print("\tType: {}".format(result['value_type']))
            print("\tData: {}".format(result['value_data']))
            print()

        elif args.get_task:
            return hyper_lr.get_scheduled_tasks()

        else:
            # perform full live response collection
            if config['lerc_install_cmd']:
                result = hyper_lr.get_lerc_status()
                if not result or result == 'UNINSTALLED' or result == 'UNKNOWN':
                   if not hyper_lr.deploy_lerc(config['lerc_install_cmd']):
                       LOGGER.warn("LERC deployment failed")
            else:
                LOGGER.info("{} environment is not configrued for LERC deployment".format(profile))
            return LR_collection(hyper_lr, args)

    # Remediation #
    if args.command == 'remediate':
        return Remediation(sensor, args)


    # Process Investigation #
    process_tree = None
    if args.command == 'proc':
        proc = proc_search_environments(profiles, args.process)
        if not proc:
            return 1
        sp = SuperProcess(proc)
        if args.proc_info:
            print(sp)
        elif args.walk_tree:
            sp.walk_process_tree()
            print()
            print(sp.process_tree)

            for process in sp.process_tree:
                if process.is_suppressed:
                    print("+  [DATA SUPPRESSED] {} (PID:{}) - {}".format(process.name, process.pid,
                                                                         process.id))
                    continue

                print("+  {} (PID:{}) - {}".format(process.name, process.pid, process.id))
                if args.filemods:
                    process.print_filemods()
                    args.no_analysis = True
                if args.netconns:
                    process.print_netconns()
                    args.no_analysis = True
                if args.regmods:
                    process.print_regmods()
                    args.no_analysis = True
                if args.unsigned_modloads:
                    process.print_unsigned_modloads()
                    args.no_analysis = True
                if args.modloads:
                    process.print_modloads()
                    args.no_analysis = True
                if args.crossprocs:
                    process.print_crossprocs()
                    args.no_analysis = True
                if args.walk_parents:
                    sp.show_ancestry()
                    args.no_analysis = True
                if args.no_analysis != True:
                    if args.json:
                        if args.segment_limit:
                            print(process.events_to_json(segment_limit=args.segment_limit))
                        else:
                            print(process.events_to_json())
                    else:
                        process.default_print()
        else:
            print()
            print(sp.process_tree)
            if args.walk_parents:
                sp.show_ancestry()
                args.no_analysis = True
            if args.filemods:
                sp.print_filemods()
                args.no_analysis = True
            if args.netconns:
                sp.print_netconns()
                args.no_analysis = True
            if args.regmods:
                sp.print_regmods()
                args.no_analysis = True
            if args.unsigned_modloads:
                sp.print_unsigned_modloads()
                args.no_analysis = True
            if args.modloads:
                sp.print_modloads()
                args.no_analysis = True
            if args.crossprocs:
                sp.print_crossprocs()
                args.no_analysis = True
            if args.show_children:
                sp.print_child_events()
                args.no_analysis = True

            if args.no_analysis != True:
                if args.json:
                    if args.segment_limit:
                        print(sp.events_to_json(segment_limit=args.segment_limit))
                    else:
                        print(sp.events_to_json())
                else:
                    sp.default_print()

        
    print()
    return True