This custom function implements a proprietary IntegralDefense Live Response collection package.
    'None' will be immediately returned if you do no have this package. The location of the package is
    assumed to be defined by a 'lr_path' variable in the configuration file.

    :dependencies lr_package_path: Path to the ID LR package on system cbinterface is running
    :dependencies streamline_path: Path to proprietary post-collection LR streamline package

    :returns: True on success, False on Failure, None if a requirement is missing

def LR_collection(hyper_lr, args):
    """This custom function implements a proprietary IntegralDefense Live Response collection package.
    'None' will be immediately returned if you do no have this package. The location of the package is
    assumed to be defined by a 'lr_path' variable in the configuration file.

    :dependencies lr_package_path: Path to the ID LR package on system cbinterface is running
    :dependencies streamline_path: Path to proprietary post-collection LR streamline package

    :returns: True on success, False on Failure, None if a requirement is missing
    """

    # Get configuration items
    config = ConfigParser()
    config.read(CONFIG_PATH)
    lr_analysis_path = config['ID-LR']['lr_package_path']
    if not os.path.exists(lr_analysis_path):
        LOGGER.info("LR package not defined")
        return None
    lr_filename = lr_analysis_path[lr_analysis_path.rfind('/')+1:]
    lr_dirname = lr_filename[:lr_filename.rfind('.')]
    sensor_dir = config['ID-LR']['sensor_path']
    if not sensor_dir:
        LOGGER.info("sensor_dir not defined in configuration. Using C:\\")
        sensor_dir = "C:\\" # Default
    streamline_path = config['ID-LR']['streamline_path']
    if not os.path.exists(streamline_path):
        LOGGER.info("Path to streamline.py doesn't exist.")
        return None    
    collect_cmd = config['ID-LR']['collect_cmd']
    if not collect_cmd:
        LOGGER.info("Collection command missing")
        return None

    lr_session = hyper_lr.go_live()

    def lr_cleanup(lr_session):
        # delete our LR tools
        try:
            dir_output = lr_session.list_directory(sensor_dir)
            for dir_item in dir_output:
                if 'DIRECTORY' in dir_item['attributes'] and dir_item['filename'] == lr_dirname:
                    print("[+] Found existing LR directory.. deleting")
                    command_str = "powershell.exe Remove-Item {} -Force -Recurse".format(sensor_dir + lr_dirname)
                    result = lr_session.create_process(command_str)
                    if result != b'':
                        LOGGER.warn(" ! Problem  deleting {}".format(sensor_dir + lr_dirname))
                if 'ARCHIVE' in dir_item['attributes'] and dir_item['filename'] == lr_filename:
                    print("[+] Existing LR package found.. deleting..")
                    try:
                        lr_session.delete_file(sensor_dir + dir_item['filename'])
                    except TypeError as e:
                        if 'startswith first arg must be bytes' in e: # might be fixed in newer cbapi versions
                            LOGGER.warn("Error deleting {}.. trying to move on".format(lr_filename))
        except live_response_api.LiveResponseError as e:
            if 'ERROR_PATH_NOT_FOUND' not in str(e):
                print("[ERROR] LiveResponseError: {}".format(e))
                return False

    # LR remnants already on host?
    lr_cleanup(lr_session)

    print("[+] Dropping latest LR on host..")
    filedata = None
    with open(lr_analysis_path, 'rb') as f: 
        filedata = f.read()
    try:
        lr_session.put_file(filedata, sensor_dir + lr_filename)
    except Exception as e:
        # If 'ERROR_FILE_EXISTS' in errmsg, log the error, but try to continue with existing package
        if 'ERROR_FILE_EXISTS' not in str(e):
            LOGGER.error("Unknown Error: {}".format(str(e)))
            return False
    
    # execute lr.exe to extract files
    # unzip = "C:\\lr.exe -o \'C:\\' -y"
    extract_cmd = " -o \'" + sensor_dir + "' -y"
    unzip = sensor_dir + lr_filename + extract_cmd
    print("[+] Extracting LR ..")
    lr_session.create_process(unzip)
    
    # execute collection
    #collect = "C:\\lr\\win32\\tools\\collect.bat"
    collect = sensor_dir + lr_dirname + collect_cmd
    collect_filename = collect_cmd[collect_cmd.rfind('\\')+1:]
    time.sleep(1)
    print("[+] Executing collect.bat..")
    start_time = time.time()
    lr_session.create_process(collect, wait_for_output=False, wait_for_completion=False) #, wait_timeout=900)
    hyper_lr.wait_for_process_to_finish(collect_filename)
    collect_duration = datetime.timedelta(minutes=(time.time() - start_time))
    print("[+] Collect completed in {}".format(collect_duration))

    # Collect resulting output file
    outputdir = sensor_dir + lr_dirname + "\\win32\\output\\"
    localfile = None
    for dir_item in lr_session.list_directory(outputdir):
        if 'ARCHIVE' in dir_item['attributes'] and  dir_item['filename'].endswith('7z'):
            # use lerc, if available
            if hyper_lr.lerc_session:
                lerc = hyper_lr.lerc_session.get_host(hyper_lr.hostname)
                command = lerc.Upload(outputdir+dir_item['filename'])
                #command = hyper_lr.lerc_session.check_command()
                # wait for client to complete the command
                print(" ~ Issued upload command to lerc. Waiting for command to finish..")
                if command.wait_for_completion():
                    print(" ~ lerc command complete. Streaming LR from lerc server..")
                    command.get_results(file_path=dir_item['filename'])
                    if command:
                        print("[+] lerc command results: ")
                        print(command)
                        file_path = command.server_file_path
                        filename = file_path[file_path.rfind('/')+1:]
                        print()
                        print("[+] Wrote {}".format(dir_item['filename']))
                        localfile = dir_item['filename']
                else:
                    LOGGER.error("problem waiting for lerc client to complete command")
            else:
                localfile = hyper_lr.getFile_with_timeout(outputdir+dir_item['filename'], localfname=dir_item['filename'])

    # Delete leftovers from sensor
    lr_cleanup(lr_session)
 
    # Call steamline on the 7z lr package
    print("[+] Starting streamline on {}".format(localfile))
    args = shlex.split(streamline_path + " " + localfile)
    subprocess.call(args, stdout=subprocess.PIPE)
    print("[+] Streamline complete")
    return True