Exemple #1
0
def GetApiClient(state):
    """Returns a Timesketch API client using thread safe methods.

  This function either returns an API client that has been stored
  in the state object, or if not it will read Timesketch RC files
  to configure a Timesketch API client. If the RC file does not exist
  or is missing values questions will be asked to fully configure
  the client.

  Args:
    state (DFTimewolfState): recipe state.

  Returns:
    object: A timesketch API object (instance of TimesketchApi).
  """
    with LOCK:
        ts_client = state.GetFromCache('timesketch_client', default_value=None)
        if ts_client:
            return ts_client

        assistant = config.ConfigAssistant()
        assistant.load_config_file()

        # TODO: support user supplied passwords to decrypt token file.
        config.configure_missing_parameters(config_assistant=assistant)
        ts_client = assistant.get_client()
        assistant.save_config()

        cred_storage = crypto.CredentialStorage()
        cred_storage.save_credentials(ts_client.credentials,
                                      config_assistant=assistant)

        state.AddToCache('timesketch_client', ts_client)
        return ts_client
Exemple #2
0
def GetApiClient(state, token_password=''):
  """Returns a Timesketch API client using thread safe methods.

  This function either returns an API client that has been stored
  in the state object, or if not it will read Timesketch RC files
  to configure a Timesketch API client. If the RC file does not exist
  or is missing values questions will be asked to fully configure
  the client.

  Args:
    state (DFTimewolfState): recipe state.
    token_password (str): optional password used to decrypt the
        Timesketch credential storage. Defaults to an empty string since
        the upstream library expects a string value. An empty string means
        a password will be generated by the upstream library.

  Returns:
    object: A timesketch API object (instance of TimesketchApi).

  Raises:
    RuntimeError: If the configuration file cannot be modified.
  """
  with LOCK:
    ts_client = state.GetFromCache('timesketch_client', default_value=None)
    if ts_client:
      return ts_client

    assistant = config.ConfigAssistant()
    assistant.load_config_file()
    try:
      config.configure_missing_parameters(
          config_assistant=assistant, token_password=token_password)
    except OSError:
      state.ModuleError(
          'Unable to get a Timesketch API Client', critical=False)
      return None


    ts_client = assistant.get_client(token_password=token_password)

    if not ts_client:
      state.ModuleError(
          'Unable to get a Timesketch API Client', critical=False)
      return None

    assistant.save_config()

    if ts_client.credentials:
      cred_storage = crypto.CredentialStorage()
      cred_storage.save_credentials(
          ts_client.credentials, config_assistant=assistant,
          password=token_password)

    state.AddToCache('timesketch_client', ts_client)
    return ts_client
Exemple #3
0
    def __init__(self, api_client=None, sketch_from_flag=None, conf_file=''):
        """Initialize the state object.

        Args:
            sketch_from_flag: Sketch ID if provided by flag.
        """
        self.api = api_client
        self.sketch_from_flag = sketch_from_flag

        if not api_client:
            try:
                # TODO: Consider other config sections here as well.
                self.api = timesketch_config.get_client(load_cli_config=True)
                if not self.api:
                    raise RequestConnectionError
            except RequestConnectionError:
                click.echo('ERROR: Cannot connect to the Timesketch server.')
                sys.exit(1)

        self.config_assistant = timesketch_config.ConfigAssistant()
        self.config_assistant.load_config_file(conf_file, load_cli_config=True)
        '--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', type=str, help=(
            'Path to the file that is to be imported.'))

    options = argument_parser.parse_args()

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

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

    cred_storage = crypto.CredentialStorage()
    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 as e:
        logger.error(
Exemple #5
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
Exemple #6
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)