Example #1
0
def main(args=None):
    """The main function of the tool."""
    if args is None:
        args = sys.argv[1:]

    argument_parser = argparse.ArgumentParser(
        description='A tool to upload data to Timesketch, using the API.')

    argument_parser.add_argument('--version',
                                 action='store_true',
                                 dest='show_version',
                                 help='Print version information')

    argument_parser.add_argument(
        '--debug',
        '--verbose',
        '-d',
        action='store_true',
        dest='show_debug',
        help='Make the logging more verbose to include debug logs.')

    auth_group = argument_parser.add_argument_group(
        title='Authentication Arguments',
        description=(
            'If no authentication parameters are supplied the default '
            'timesketch RC and token files will be used to provide the '
            'authentication information. If those files are not present '
            'the tool will ask you questions and store the results in those '
            'files for future authentication.'))

    auth_group.add_argument('-u',
                            '--user',
                            '--username',
                            action='store',
                            dest='username',
                            type=str,
                            help='The username of the Timesketch user.')
    auth_group.add_argument(
        '-p',
        '--password',
        '--pwd',
        action='store',
        type=str,
        dest='password',
        help=(
            'If authenticated with password, provide the password on the CLI. '
            'If neither password is provided nor a password prompt an OAUTH '
            'connection is assumed.'))
    auth_group.add_argument('--pwd-prompt',
                            '--pwd_prompt',
                            action='store_true',
                            default=False,
                            dest='pwd_prompt',
                            help='Prompt for password.')
    auth_group.add_argument(
        '--cred-prompt',
        '--cred_prompt',
        '--token-password',
        '--token_password',
        '--token',
        action='store_true',
        default=False,
        dest='cred_prompt',
        help='Prompt for password to decrypt and encrypt credential file.')
    auth_group.add_argument('--client-secret',
                            '--client_secret',
                            action='store',
                            type=str,
                            default='',
                            dest='client_secret',
                            help='OAUTH client secret.')
    auth_group.add_argument('--client-id',
                            '--client_id',
                            action='store',
                            type=str,
                            default='',
                            dest='client_id',
                            help='OAUTH client ID.')
    auth_group.add_argument(
        '--run_local',
        '--run-local',
        action='store_true',
        dest='run_local',
        help=('If OAUTH is used to authenticate and the connection is over '
              'SSH then it is recommended to set this option. When set an '
              'authentication URL is prompted on the screen, requiring a '
              'copy/paste into a browser to complete the OAUTH dance.'))
    auth_group.add_argument(
        '--re-prompt',
        '--re_prompt',
        action='store_true',
        dest='re_prompt',
        help=('If for some reasons you type the wrong username, password '
              'or host information you can use this parameter to re-enter '
              'those mistyped values.'))

    config_group = argument_parser.add_argument_group(
        'Configuration Arguments')

    config_group.add_argument(
        '--quick',
        '-q',
        '--no-wait',
        '--no_wait',
        action='store_false',
        default=True,
        dest='wait_timeline',
        help=('By default the tool will wait until the timeline has been '
              'indexed and print out some details of the import. This option '
              'makes the tool exit as soon as the data has been imported and '
              'does not wait until it\'s been indexed.'))

    config_group.add_argument(
        '--log-config-file',
        '--log_config_file',
        '--lc',
        action='store',
        type=str,
        default='',
        metavar='FILEPATH',
        dest='log_config_file',
        help=('Path to a YAML config file that defines the config for parsing '
              'and setting up file parsing. By default formatter.yaml that '
              'comes with the importer will be used.'))

    config_group.add_argument('--host',
                              '--hostname',
                              '--host-uri',
                              '--host_uri',
                              dest='host_uri',
                              type=str,
                              default='',
                              action='store',
                              help='The URI to the Timesketch instance')

    config_group.add_argument(
        '--format_string',
        '--format-string',
        type=str,
        action='store',
        dest='format_string',
        default='',
        help=(
            'Formatting string for the message field. If there is no message '
            'field in the input data a message string can be composed using '
            'a format string.'))

    config_group.add_argument(
        '--timeline_name',
        '--timeline-name',
        action='store',
        type=str,
        dest='timeline_name',
        default='',
        help=('String that will be used as the timeline name.'))

    config_group.add_argument(
        '--sketch_name',
        '--sketch-name',
        action='store',
        type=str,
        dest='sketch_name',
        default='',
        help=('String that will be used as the sketch name in case a new '
              'sketch is created.'))

    config_group.add_argument(
        '--data_label',
        '--data-label',
        action='store',
        type=str,
        dest='data_label',
        default='',
        help=(
            'The data label is used by the API to determine whether a new '
            'search index needs to be created or if the data can be appended '
            'to an already existing index. If a file is added this defaults '
            'to the file extension, otherwise a default value of generic is '
            'applied.'))

    config_group.add_argument(
        '--config_section',
        '--config-section',
        action='store',
        type=str,
        dest='config_section',
        default='',
        help=('The config section in the RC file that will be used to '
              'define server information.'))

    config_group.add_argument(
        '--index-name',
        '--index_name',
        action='store',
        type=str,
        default='',
        dest='index_name',
        help=('If the data should be imported into a specific timeline the '
              'index name needs to be provided, otherwise a new index will '
              'be generated.'))
    config_group.add_argument(
        '--timestamp_description',
        '--timestamp-description',
        '--time-desc',
        '--time_desc',
        action='store',
        type=str,
        default='',
        dest='time_desc',
        help='Value for the timestamp_description field.')

    config_group.add_argument(
        '--threshold_entry',
        '--threshold-entry',
        '--entries',
        action='store',
        type=int,
        default=0,
        dest='entry_threshold',
        help=('How many entries should be buffered up before being '
              'sent to server.'))

    config_group.add_argument(
        '--threshold_size',
        '--threshold-size',
        '--filesize',
        action='store',
        type=int,
        default=0,
        dest='size_threshold',
        help=('For binary file transfer, how many bytes should be transferred '
              'per chunk.'))

    config_group.add_argument(
        '--sketch_id',
        '--sketch-id',
        type=int,
        default=0,
        dest='sketch_id',
        action='store',
        help=('The sketch ID to store the timeline in, if no sketch ID is '
              'provided a new sketch will be created.'))

    config_group.add_argument(
        '--context',
        action='store',
        type=str,
        default='',
        dest='context',
        help=('Set a context for the file upload. This could be a text '
              'describing how the data got collected or parameters to '
              'the tool. Defaults to how the CLI tool got run.'))

    argument_parser.add_argument(
        'path',
        action='store',
        nargs='?',
        type=str,
        help=('Path to the file that is to be imported.'))

    options = argument_parser.parse_args(args)

    if options.show_version:
        print('API Client Version: {0:s}'.format(api_version.get_version()))
        print('Importer Client Version: {0:s}'.format(
            importer_version.get_version()))
        sys.exit(0)

    if options.show_debug:
        configure_logger_debug()
    else:
        configure_logger_default()

    if not options.path:
        logger.error(
            'A valid file path needs to be provided, unable to continue.')
        sys.exit(1)

    if not os.path.isfile(options.path):
        logger.error('Path {0:s} is not valid, unable to continue.'.format(
            options.path))
        sys.exit(1)

    config_section = options.config_section
    assistant = config.ConfigAssistant()
    assistant.load_config_file(section=config_section)
    assistant.load_config_dict(vars(options))

    try:
        file_path = assistant.get_config('token_file_path')
    except KeyError:
        file_path = ''

    cred_storage = crypto.CredentialStorage(file_path=file_path)
    token_password = ''
    if options.cred_prompt:
        token_password = cli_input.ask_question(
            'Enter password to encrypt/decrypt credential file',
            input_type=str,
            hide_input=True)

    try:
        credentials = cred_storage.load_credentials(config_assistant=assistant,
                                                    password=token_password)
    except IOError:
        logger.error('Unable to decrypt the credential file')
        logger.debug('Error details:', exc_info=True)
        logger.error('If you\'ve forgotten the password you can delete '
                     'the ~/.timesketch.token file and run the tool again.')
        sys.exit(1)

    conf_password = ''

    if credentials:
        logger.info('Using cached credentials.')
        if credentials.TYPE.lower() == 'oauth':
            assistant.set_config('auth_mode', 'oauth')
        elif credentials.TYPE.lower() == 'timesketch':
            assistant.set_config('auth_mode', 'timesketch')

    else:
        # Check whether we'll need to read password for timesketch
        # auth from the command line.
        if options.pwd_prompt:
            conf_password = getpass.getpass('Type in the password: '******'auth_mode', 'oauth_local')

        if options.client_secret:
            assistant.set_config('auth_mode', 'oauth')

        if conf_password:
            assistant.set_config('auth_mode', 'timesketch')

    if conf_password:
        credentials = ts_credentials.TimesketchPwdCredentials()
        credentials.credential = {
            'username': assistant.get_config('username'),
            'password': conf_password
        }
        logger.info('Saving Credentials.')
        cred_storage.save_credentials(credentials,
                                      password=token_password,
                                      file_path=file_path,
                                      config_assistant=assistant)

    # Gather all questions that are missing.
    config.configure_missing_parameters(config_assistant=assistant,
                                        token_password=token_password,
                                        confirm_choices=options.re_prompt)

    logger.info('Creating a client.')
    ts_client = assistant.get_client(token_password=token_password)
    if not ts_client:
        logger.error('Unable to create a Timesketch API client, exiting.')
        sys.exit(1)

    logger.info('Client created.')
    logger.info('Saving TS config.')
    assistant.save_config(section=config_section, token_file_path=file_path)

    if ts_client.credentials:
        logger.info('Saving Credentials.')
        cred_storage.save_credentials(ts_client.credentials,
                                      password=token_password,
                                      file_path=file_path,
                                      config_assistant=assistant)

    sketch_id = options.sketch_id
    if sketch_id:
        my_sketch = ts_client.get_sketch(sketch_id)
    else:
        sketch_name = options.sketch_name or 'New Sketch From Importer CLI'
        my_sketch = ts_client.create_sketch(sketch_name)
        logger.info('New sketch created: [{0:d}] {1:s}'.format(
            my_sketch.id, my_sketch.name))

    if not my_sketch:
        logger.error('Unable to get sketch ID: {0:d}'.format(sketch_id))
        sys.exit(1)

    filename = os.path.basename(options.path)
    default_timeline_name, _, _ = filename.rpartition('.')

    if options.timeline_name:
        conf_timeline_name = options.timeline_name
    else:
        conf_timeline_name = cli_input.ask_question(
            'What is the timeline name',
            input_type=str,
            default=default_timeline_name)

    config_dict = {
        'message_format_string': options.format_string,
        'timeline_name': conf_timeline_name,
        'index_name': options.index_name,
        'timestamp_description': options.time_desc,
        'entry_threshold': options.entry_threshold,
        'size_threshold': options.size_threshold,
        'log_config_file': options.log_config_file,
        'data_label': options.data_label,
        'context': options.context,
    }

    logger.info('Uploading file.')
    timeline, task_id = upload_file(my_sketch=my_sketch,
                                    config_dict=config_dict,
                                    file_path=options.path)

    if not options.wait_timeline:
        logger.info('File got successfully uploaded to sketch: {0:d}'.format(
            my_sketch.id))
        return

    if not timeline:
        logger.warning(
            'There does not seem to be any timeline returned, check whether '
            'the data got uploaded.')
        return

    print('Checking file upload status: ', end='')
    while True:
        status = timeline.status
        if status in ('archived', 'failed', 'fail'):
            print('[FAIL]')
            print('Unable to index timeline {0:s}, reason: {1:s}'.format(
                timeline.description, timeline.status))
            return

        if status not in ('ready', 'success'):
            print('.', end='')
            time.sleep(3)
            continue

        print('[DONE]')
        print(f'Timeline uploaded to Timeline Id: {timeline.id}.')

        task_state = 'Unknown'
        task_list = ts_client.check_celery_status(task_id)
        for task in task_list:
            if task.get('task_id', '') == task_id:
                task_state = task.get('state', 'Unknown')
        print(f'Status of the index is: {task_state}')
        break
Example #2
0
from timesketch_api_client import version

long_description = (
    'The Timesketch API client provides you with a set of Python libraries '
    'to connect to your Timesketch (https://github.com/google/timesketch) '
    'instance.\n\n'
    'The API is feature complete with the Timesketch UI and allows you to '
    'do all operations that can be done in the UI, providing ways to '
    'integrate Timesketch into other products such as Jupyter/Colab.\n\n'
    'To see how it works in action, try the colab notebook that is accessible '
    'from here: https://colab.research.google.com/github/google/timesketch/'
    'blob/master/notebooks/colab-timesketch-demo.ipynb')

setup(
    name='timesketch-api-client',
    version=version.get_version(),
    description='Timesketch API client',
    long_description=long_description,
    license='Apache License, Version 2.0',
    url='http://www.timesketch.org/',
    maintainer='Timesketch development team',
    maintainer_email='*****@*****.**',
    classifiers=[
        'Development Status :: 4 - Beta',
        'Environment :: Console',
        'Operating System :: OS Independent',
        'Programming Language :: Python',
    ],
    packages=find_packages(),
    include_package_data=True,
    zip_safe=False,
Example #3
0
def main(args=None):
    """The main function of the tool."""
    if args is None:
        args = sys.argv[1:]

    argument_parser = argparse.ArgumentParser(
        description="A tool to upload data to Timesketch, using the API.")

    argument_parser.add_argument(
        "--version",
        action="store_true",
        dest="show_version",
        help="Print version information",
    )

    argument_parser.add_argument(
        "--debug",
        "--verbose",
        "-d",
        action="store_true",
        dest="show_debug",
        help="Make the logging more verbose to include debug logs.",
    )

    auth_group = argument_parser.add_argument_group(
        title="Authentication Arguments",
        description=(
            "If no authentication parameters are supplied the default "
            "timesketch RC and token files will be used to provide the "
            "authentication information. If those files are not present "
            "the tool will ask you questions and store the results in those "
            "files for future authentication."),
    )

    auth_group.add_argument(
        "-u",
        "--user",
        "--username",
        action="store",
        dest="username",
        type=str,
        help="The username of the Timesketch user.",
    )
    auth_group.add_argument(
        "-p",
        "--password",
        "--pwd",
        action="store",
        type=str,
        dest="password",
        help=(
            "If authenticated with password, provide the password on the CLI. "
            "If neither password is provided nor a password prompt an OAUTH "
            "connection is assumed."),
    )
    auth_group.add_argument(
        "--pwd-prompt",
        "--pwd_prompt",
        action="store_true",
        default=False,
        dest="pwd_prompt",
        help="Prompt for password.",
    )
    auth_group.add_argument(
        "--cred-prompt",
        "--cred_prompt",
        "--token-password",
        "--token_password",
        "--token",
        action="store_true",
        default=False,
        dest="cred_prompt",
        help="Prompt for password to decrypt and encrypt credential file.",
    )
    auth_group.add_argument(
        "--client-secret",
        "--client_secret",
        action="store",
        type=str,
        default="",
        dest="client_secret",
        help="OAUTH client secret.",
    )
    auth_group.add_argument(
        "--client-id",
        "--client_id",
        action="store",
        type=str,
        default="",
        dest="client_id",
        help="OAUTH client ID.",
    )
    auth_group.add_argument(
        "--run_local",
        "--run-local",
        action="store_true",
        dest="run_local",
        help=("If OAUTH is used to authenticate and the connection is over "
              "SSH then it is recommended to set this option. When set an "
              "authentication URL is prompted on the screen, requiring a "
              "copy/paste into a browser to complete the OAUTH dance."),
    )
    auth_group.add_argument(
        "--re-prompt",
        "--re_prompt",
        action="store_true",
        dest="re_prompt",
        help=("If for some reasons you type the wrong username, password "
              "or host information you can use this parameter to re-enter "
              "those mistyped values."),
    )

    config_group = argument_parser.add_argument_group(
        "Configuration Arguments")

    config_group.add_argument(
        "--quick",
        "-q",
        "--no-wait",
        "--no_wait",
        action="store_false",
        default=True,
        dest="wait_timeline",
        help=("By default the tool will wait until the timeline has been "
              "indexed and print out some details of the import. This option "
              "makes the tool exit as soon as the data has been imported and "
              "does not wait until it's been indexed."),
    )

    config_group.add_argument(
        "--log-config-file",
        "--log_config_file",
        "--lc",
        action="store",
        type=str,
        default="",
        metavar="FILEPATH",
        dest="log_config_file",
        help=("Path to a YAML config file that defines the config for parsing "
              "and setting up file parsing. By default formatter.yaml that "
              "comes with the importer will be used."),
    )

    config_group.add_argument(
        "--host",
        "--hostname",
        "--host-uri",
        "--host_uri",
        dest="host_uri",
        type=str,
        default="",
        action="store",
        help="The URI to the Timesketch instance",
    )

    config_group.add_argument(
        "--format_string",
        "--format-string",
        type=str,
        action="store",
        dest="format_string",
        default="",
        help=(
            "Formatting string for the message field. If there is no message "
            "field in the input data a message string can be composed using "
            "a format string."),
    )

    config_group.add_argument(
        "--timeline_name",
        "--timeline-name",
        action="store",
        type=str,
        dest="timeline_name",
        default="",
        help=("String that will be used as the timeline name."),
    )

    config_group.add_argument(
        "--sketch_name",
        "--sketch-name",
        action="store",
        type=str,
        dest="sketch_name",
        default="",
        help=("String that will be used as the sketch name in case a new "
              "sketch is created."),
    )

    config_group.add_argument(
        "--data_label",
        "--data-label",
        action="store",
        type=str,
        dest="data_label",
        default="",
        help=(
            "The data label is used by the API to determine whether a new "
            "search index needs to be created or if the data can be appended "
            "to an already existing index. If a file is added this defaults "
            "to the file extension, otherwise a default value of generic is "
            "applied."),
    )

    config_group.add_argument(
        "--config_section",
        "--config-section",
        action="store",
        type=str,
        dest="config_section",
        default="",
        help=("The config section in the RC file that will be used to "
              "define server information."),
    )

    config_group.add_argument(
        "--index-name",
        "--index_name",
        action="store",
        type=str,
        default="",
        dest="index_name",
        help=("If the data should be imported into a specific timeline the "
              "index name needs to be provided, otherwise a new index will "
              "be generated."),
    )
    config_group.add_argument(
        "--timestamp_description",
        "--timestamp-description",
        "--time-desc",
        "--time_desc",
        action="store",
        type=str,
        default="",
        dest="time_desc",
        help="Value for the timestamp_description field.",
    )

    config_group.add_argument(
        "--threshold_entry",
        "--threshold-entry",
        "--entries",
        action="store",
        type=int,
        default=0,
        dest="entry_threshold",
        help=("How many entries should be buffered up before being "
              "sent to server."),
    )

    config_group.add_argument(
        "--threshold_size",
        "--threshold-size",
        "--filesize",
        action="store",
        type=int,
        default=0,
        dest="size_threshold",
        help=("For binary file transfer, how many bytes should be transferred "
              "per chunk."),
    )

    config_group.add_argument(
        "--sketch_id",
        "--sketch-id",
        type=int,
        default=0,
        dest="sketch_id",
        action="store",
        help=("The sketch ID to store the timeline in, if no sketch ID is "
              "provided a new sketch will be created."),
    )

    config_group.add_argument(
        "--context",
        action="store",
        type=str,
        default="",
        dest="context",
        help=("Set a context for the file upload. This could be a text "
              "describing how the data got collected or parameters to "
              "the tool. Defaults to how the CLI tool got run."),
    )

    argument_parser.add_argument(
        "path",
        action="store",
        nargs="?",
        type=str,
        help=("Path to the file that is to be imported."),
    )

    options = argument_parser.parse_args(args)

    if options.show_version:
        print("API Client Version: {0:s}".format(api_version.get_version()))
        print("Importer Client Version: {0:s}".format(
            importer_version.get_version()))
        sys.exit(0)

    if options.show_debug:
        configure_logger_debug()
    else:
        configure_logger_default()

    if not options.path:
        logger.error(
            "A valid file path needs to be provided, unable to continue.")
        sys.exit(1)

    if not os.path.isfile(options.path):
        logger.error("Path {0:s} is not valid, unable to continue.".format(
            options.path))
        sys.exit(1)

    config_section = options.config_section
    assistant = config.ConfigAssistant()
    assistant.load_config_file(section=config_section)
    assistant.load_config_dict(vars(options))

    try:
        file_path = assistant.get_config("token_file_path")
    except KeyError:
        file_path = ""

    cred_storage = crypto.CredentialStorage(file_path=file_path)
    token_password = ""
    if options.cred_prompt:
        token_password = cli_input.ask_question(
            "Enter password to encrypt/decrypt credential file",
            input_type=str,
            hide_input=True,
        )

    try:
        credentials = cred_storage.load_credentials(config_assistant=assistant,
                                                    password=token_password)
    except IOError:
        logger.error("Unable to decrypt the credential file")
        logger.debug("Error details:", exc_info=True)
        logger.error("If you've forgotten the password you can delete "
                     "the ~/.timesketch.token file and run the tool again.")
        sys.exit(1)

    conf_password = ""

    if credentials:
        logger.info("Using cached credentials.")
        if credentials.TYPE.lower() == "oauth":
            assistant.set_config("auth_mode", "oauth")
        elif credentials.TYPE.lower() == "timesketch":
            assistant.set_config("auth_mode", "timesketch")

    else:
        # Check whether we'll need to read password for timesketch
        # auth from the command line.
        if options.pwd_prompt:
            conf_password = getpass.getpass("Type in the password: "******"auth_mode", "oauth_local")

        if options.client_secret:
            assistant.set_config("auth_mode", "oauth")

        if conf_password:
            assistant.set_config("auth_mode", "timesketch")

    if conf_password:
        credentials = ts_credentials.TimesketchPwdCredentials()
        credentials.credential = {
            "username": assistant.get_config("username"),
            "password": conf_password,
        }
        logger.info("Saving Credentials.")
        cred_storage.save_credentials(
            credentials,
            password=token_password,
            file_path=file_path,
            config_assistant=assistant,
        )

    # Gather all questions that are missing.
    config.configure_missing_parameters(
        config_assistant=assistant,
        token_password=token_password,
        confirm_choices=options.re_prompt,
    )

    logger.info("Creating a client.")
    ts_client = assistant.get_client(token_password=token_password)
    if not ts_client:
        logger.error("Unable to create a Timesketch API client, exiting.")
        sys.exit(1)

    logger.info("Client created.")
    logger.info("Saving TS config.")
    assistant.save_config(section=config_section, token_file_path=file_path)

    if ts_client.credentials:
        logger.info("Saving Credentials.")
        cred_storage.save_credentials(
            ts_client.credentials,
            password=token_password,
            file_path=file_path,
            config_assistant=assistant,
        )

    sketch_id = options.sketch_id
    if sketch_id:
        my_sketch = ts_client.get_sketch(sketch_id)
    else:
        sketch_name = options.sketch_name or "New Sketch From Importer CLI"
        my_sketch = ts_client.create_sketch(sketch_name)
        logger.info("New sketch created: [{0:d}] {1:s}".format(
            my_sketch.id, my_sketch.name))

    if not my_sketch:
        logger.error("Unable to get sketch ID: {0:d}".format(sketch_id))
        sys.exit(1)

    filename = os.path.basename(options.path)
    default_timeline_name, _, _ = filename.rpartition(".")

    if options.timeline_name:
        conf_timeline_name = options.timeline_name
    else:
        conf_timeline_name = cli_input.ask_question(
            "What is the timeline name",
            input_type=str,
            default=default_timeline_name)

    config_dict = {
        "message_format_string": options.format_string,
        "timeline_name": conf_timeline_name,
        "index_name": options.index_name,
        "timestamp_description": options.time_desc,
        "entry_threshold": options.entry_threshold,
        "size_threshold": options.size_threshold,
        "log_config_file": options.log_config_file,
        "data_label": options.data_label,
        "context": options.context,
    }

    logger.info("Uploading file.")
    timeline, task_id = upload_file(my_sketch=my_sketch,
                                    config_dict=config_dict,
                                    file_path=options.path)

    if not options.wait_timeline:
        logger.info("File got successfully uploaded to sketch: {0:d}".format(
            my_sketch.id))
        return

    if not timeline:
        logger.warning(
            "There does not seem to be any timeline returned, check whether "
            "the data got uploaded.")
        return

    print("Checking file upload status: ", end="")
    while True:
        status = timeline.status
        if status in ("archived", "failed", "fail"):
            print("[FAIL]")
            print("Unable to index timeline {0:s}, reason: {1:s}".format(
                timeline.description, timeline.status))
            return

        if status not in ("ready", "success"):
            print(".", end="")
            time.sleep(3)
            continue

        print("[DONE]")
        print(f"Timeline uploaded to Timeline Id: {timeline.id}.")

        task_state = "Unknown"
        task_list = ts_client.check_celery_status(task_id)
        for task in task_list:
            if task.get("task_id", "") == task_id:
                task_state = task.get("state", "Unknown")
        print(f"Status of the index is: {task_state}")
        break
def main(args=None):
    """The main function of the tool."""
    if args is None:
        args = sys.argv[1:]

    argument_parser = argparse.ArgumentParser(
        description='A tool to upload data to Timesketch, using the API.')

    argument_parser.add_argument('--version',
                                 action='store_true',
                                 dest='show_version',
                                 help='Print version information')

    argument_parser.add_argument(
        '--debug',
        '--verbose',
        '-d',
        action='store_true',
        dest='show_debug',
        help='Make the logging more verbose to include debug logs.')

    auth_group = argument_parser.add_argument_group('Authentication Arguments')
    auth_group.add_argument('-u',
                            '--user',
                            '--username',
                            action='store',
                            dest='username',
                            type=str,
                            help='The username of the Timesketch user.')
    auth_group.add_argument(
        '-p',
        '--password',
        '--pwd',
        action='store',
        type=str,
        dest='password',
        help=(
            'If authenticated with password, provide the password on the CLI. '
            'If neither password is provided nor a password prompt an OAUTH '
            'connection is assumed.'))
    auth_group.add_argument('--pwd-prompt',
                            '--pwd_prompt',
                            action='store_true',
                            default=False,
                            dest='pwd_prompt',
                            help='Prompt for password.')
    auth_group.add_argument(
        '--cred-prompt',
        '--cred_prompt',
        '--token-password',
        '--token_password',
        '--token',
        action='store_true',
        default=False,
        dest='cred_prompt',
        help='Prompt for password to decrypt and encrypt credential file.')
    auth_group.add_argument('--client-secret',
                            '--client_secret',
                            action='store',
                            type=str,
                            default='',
                            dest='client_secret',
                            help='OAUTH client secret.')
    auth_group.add_argument('--client-id',
                            '--client_id',
                            action='store',
                            type=str,
                            default='',
                            dest='client_id',
                            help='OAUTH client ID.')
    auth_group.add_argument(
        '--run_local',
        '--run-local',
        action='store_true',
        dest='run_local',
        help=('If OAUTH is used to authenticate and the connection is over '
              'SSH then it is recommended to set this option. When set an '
              'authentication URL is prompted on the screen, requiring a '
              'copy/paste into a browser to complete the OAUTH dance.'))

    config_group = argument_parser.add_argument_group(
        'Configuration Arguments')
    config_group.add_argument(
        '--log-config-file',
        '--log_config_file',
        '--lc',
        action='store',
        type=str,
        default='',
        metavar='FILEPATH',
        dest='log_config_file',
        help=('Path to a YAML config file that defines the config for parsing '
              'and setting up file parsing. By default formatter.yaml that '
              'comes with the importer will be used.'))
    config_group.add_argument('--host',
                              '--hostname',
                              '--host-uri',
                              '--host_uri',
                              dest='host_uri',
                              type=str,
                              default='',
                              action='store',
                              help='The URI to the Timesketch instance')
    config_group.add_argument(
        '--format_string',
        '--format-string',
        type=str,
        action='store',
        dest='format_string',
        default='',
        help=(
            'Formatting string for the message field. If there is no message '
            'field in the input data a message string can be composed using '
            'a format string.'))
    config_group.add_argument(
        '--timeline_name',
        '--timeline-name',
        action='store',
        type=str,
        dest='timeline_name',
        default='',
        help=('String that will be used as the timeline name.'))
    config_group.add_argument(
        '--config_section',
        '--config-section',
        action='store',
        type=str,
        dest='config_section',
        default='',
        help=('The config section in the RC file that will be used to '
              'define server information.'))
    config_group.add_argument(
        '--index-name',
        '--index_name',
        action='store',
        type=str,
        default='',
        dest='index_name',
        help=('If the data should be imported into a specific timeline the '
              'index name needs to be provided, otherwise a new index will '
              'be generated.'))
    config_group.add_argument(
        '--timestamp_description',
        '--timestamp-description',
        '--time-desc',
        '--time_desc',
        action='store',
        type=str,
        default='',
        dest='time_desc',
        help='Value for the timestamp_description field.')
    config_group.add_argument(
        '--threshold_entry',
        '--threshold-entry',
        '--entries',
        action='store',
        type=int,
        default=0,
        dest='entry_threshold',
        help=('How many entries should be buffered up before being '
              'sent to server.'))
    config_group.add_argument(
        '--threshold_size',
        '--threshold-size',
        '--filesize',
        action='store',
        type=int,
        default=0,
        dest='size_threshold',
        help=('For binary file transfer, how many bytes should be transferred '
              'per chunk.'))
    config_group.add_argument(
        '--sketch_id',
        '--sketch-id',
        type=int,
        default=0,
        dest='sketch_id',
        action='store',
        help=('The sketch ID to store the timeline in, if no sketch ID is '
              'provided a new sketch will be created.'))

    argument_parser.add_argument(
        'path',
        action='store',
        nargs='?',
        type=str,
        help=('Path to the file that is to be imported.'))

    options = argument_parser.parse_args(args)

    if options.show_version:
        print('API Client Version: {0:s}'.format(api_version.get_version()))
        print('Importer Client Version: {0:s}'.format(
            importer_version.get_version()))
        sys.exit(0)

    if options.show_debug:
        configure_logger_debug()
    else:
        configure_logger_default()

    if not options.path:
        logger.error(
            'A valid file path needs to be provided, unable to continue.')
        sys.exit(1)

    if not os.path.isfile(options.path):
        logger.error('Path {0:s} is not valid, unable to continue.'.format(
            options.path))
        sys.exit(1)

    config_section = options.config_section
    assistant = config.ConfigAssistant()
    assistant.load_config_file(section=config_section)
    assistant.load_config_dict(vars(options))

    try:
        file_path = assistant.get_config('token_file_path')
    except KeyError:
        file_path = ''

    cred_storage = crypto.CredentialStorage(file_path=file_path)
    token_password = ''
    if options.cred_prompt:
        token_password = cli_input.ask_question(
            'Enter password to encrypt/decrypt credential file',
            input_type=str,
            hide_input=True)

    try:
        credentials = cred_storage.load_credentials(config_assistant=assistant,
                                                    password=token_password)
    except IOError:
        logger.error('Unable to decrypt the credential file')
        logger.debug('Error details:', exc_info=True)
        logger.error('If you\'ve forgotten the password you can delete '
                     'the ~/.timesketch.token file and run the tool again.')
        sys.exit(1)

    conf_password = ''

    if credentials:
        if credentials.TYPE.lower() == 'oauth':
            assistant.set_config('auth_mode', 'oauth')
        elif credentials.TYPE.lower() == 'timesketch':
            assistant.set_config('auth_mode', 'timesketch')

    else:
        # Check whether we'll need to read password for timesketch
        # auth from the command line.
        if options.pwd_prompt:
            conf_password = getpass.getpass('Type in the password: '******'auth_mode', 'oauth_local')

        if options.client_secret:
            assistant.set_config('auth_mode', 'oauth')

        if conf_password:
            assistant.set_config('auth_mode', 'timesketch')

    if conf_password:
        credentials = ts_credentials.TimesketchPwdCredentials()
        credentials.credential = {
            'username': assistant.get_config('username'),
            'password': conf_password
        }
        logger.info('Saving Credentials.')
        cred_storage.save_credentials(credentials,
                                      password=token_password,
                                      file_path=file_path,
                                      config_assistant=assistant)

    # Gather all questions that are missing.
    config.configure_missing_parameters(config_assistant=assistant,
                                        token_password=token_password)

    logger.info('Creating a client.')
    ts_client = assistant.get_client(token_password=token_password)
    if not ts_client:
        logger.error('Unable to create a Timesketch API client, exiting.')
        sys.exit(1)

    logger.info('Client created.')
    logger.info('Saving TS config.')
    assistant.save_config(section=config_section, token_file_path=file_path)

    if ts_client.credentials:
        logger.info('Saving Credentials.')
        cred_storage.save_credentials(ts_client.credentials,
                                      password=token_password,
                                      file_path=file_path,
                                      config_assistant=assistant)

    sketch_id = options.sketch_id
    if sketch_id:
        my_sketch = ts_client.get_sketch(sketch_id)
    else:
        my_sketch = ts_client.create_sketch('New Sketch From Importer CLI')

    if not my_sketch:
        logger.error('Unable to get sketch ID: {0:d}'.format(sketch_id))
        sys.exit(1)

    filename = os.path.basename(options.path)
    default_timeline_name, _, _ = filename.rpartition('.')

    if options.timeline_name:
        conf_timeline_name = options.timeline_name
    else:
        conf_timeline_name = cli_input.ask_question(
            'What is the timeline name',
            input_type=str,
            default=default_timeline_name)

    config_dict = {
        'message_format_string': options.format_string,
        'timeline_name': conf_timeline_name,
        'index_name': options.index_name,
        'timestamp_description': options.time_desc,
        'entry_threshold': options.entry_threshold,
        'size_threshold': options.size_threshold,
        'log_config_file': options.log_config_file,
    }

    logger.info('Uploading file.')
    result = upload_file(my_sketch=my_sketch,
                         config_dict=config_dict,
                         file_path=options.path)
    logger.info(result)