diff --git a/api_client/python/timesketch_api_client/client.py b/api_client/python/timesketch_api_client/client.py index bd28161a45..d4dd127202 100644 --- a/api_client/python/timesketch_api_client/client.py +++ b/api_client/python/timesketch_api_client/client.py @@ -638,6 +638,29 @@ def get_sketch(self, sketch_id): """ return sketch.Sketch(sketch_id, api=self) + def get_sketches_by_name(self, sketch_name) -> list[sketch.Sketch]: + """Get a sketch by name. + + Args: + sketch_name (str): The name of the sketch to find. + + Raises: + KeyError: If no sketch with the specified name is found. + + Returns: + list[sketch.Sketch]: A list of sketch objects. + """ + sketches = [ + sketch_obj + for sketch_obj in self.list_sketches() + if sketch_obj.name == sketch_name + ] + + if not sketches: + raise KeyError(f"Sketch with name '{sketch_name}' not found.") + + return sketches + def get_aggregator_info(self, name="", as_pandas=False): """Returns information about available aggregators. diff --git a/api_client/python/timesketch_api_client/sketch.py b/api_client/python/timesketch_api_client/sketch.py index 2144968a95..bced8cf119 100644 --- a/api_client/python/timesketch_api_client/sketch.py +++ b/api_client/python/timesketch_api_client/sketch.py @@ -161,6 +161,42 @@ def last_activity(self): meta = data.get("meta", {}) return meta.get("last_activity", "") + @property + def created_at(self): + """Property that returns sketch creation time. + + Returns: + str: Sketch creation time as string. + """ + data = self.lazyload_data(refresh_cache=True) + objects = data.get("objects") + if not objects: + return "" + data_object = objects[0] + created_at_string = data_object.get("created_at", "") + return created_at_string + + @property + def creator(self): + """Property that returns sketch creator. + + Returns: + str: Sketch creator as string. + """ + data = self.lazyload_data(refresh_cache=True) + objects = data.get("objects") + if not objects: + return "" + data_object = objects[0] + try: + username_string = data_object.get("user").get("username") + if not username_string: + return "" + return username_string + except Exception as e: + logger.error("Error getting sketch creator: %s", e) + return "" + @property def my_acl(self): """Property that returns back the ACL for the current user.""" diff --git a/importer_client/python/tools/timesketch_importer.py b/importer_client/python/tools/timesketch_importer.py index 10952f2d6d..1c3a790848 100644 --- a/importer_client/python/tools/timesketch_importer.py +++ b/importer_client/python/tools/timesketch_importer.py @@ -26,7 +26,7 @@ from timesketch_api_client import cli_input from timesketch_api_client import credentials as ts_credentials from timesketch_api_client import crypto -from timesketch_api_client import config +from timesketch_api_client import config, client from timesketch_api_client import sketch from timesketch_api_client import version as api_version from timesketch_import_client import helper @@ -349,6 +349,22 @@ def main(args=None): ), ) + config_group.add_argument( + "--sketch_strategy", + "--sketch-strategy", + action="store", + type=str, + dest="sketch_strategy", + default="ask", + help=( + "Strategy to use when a sketch name is provided and a sketch " + "with the same name already exists. Supported strategies are: " + "'ask' (default) which will raise an error, 'newest' which will " + "use the most recently created sketch, and 'oldest' which will " + "use the earliest created sketch." + ), + ) + config_group.add_argument( "--data_label", "--data-label", @@ -601,10 +617,51 @@ def main(args=None): 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) - ) + try: + sketches = ts_client.get_sketches_by_name(sketch_name) + if len(sketches) > 1: + if options.sketch_strategy == 'newest': + my_sketch = sorted( + sketches, key=lambda s: s.created_at, reverse=True + )[0] + elif options.sketch_strategy == 'oldest': + my_sketch = sorted( + sketches, key=lambda s: s.created_at + )[0] + else: + # ask user for clarification using cli_input + print("Multiple sketches found with the name '{0:s}':".format(sketch_name)) + for s in sketches: + print(" - [{0:d}] created_at: {1:s}, by {2:s}".format( + s.id, s.created_at, s.creator)) + + selected_option = cli_input.ask_question( + "Select the sketch to use by entering the corresponding number", + input_type=int, + default=sketches[0].id, + ) + try: + # select sketch by ID from the sketches list + my_sketch = next(s for s in sketches if s.id == selected_option) + except StopIteration: + logger.error("Selected sketch ID not found, exiting.") + sys.exit(1) + else: + my_sketch = sketches[0] + + logger.info( + "Using existing sketch: [%d] %s", + my_sketch.id, + my_sketch.name, + ) + except KeyError: + # no existing sketch found, create a new one + my_sketch = ts_client.create_sketch(sketch_name) + logger.info( + "New sketch created: [%d] %s", + my_sketch.id, + my_sketch.name, + ) if not my_sketch: logger.error("Unable to get sketch ID: {0:d}".format(sketch_id))