From 9988b1c2b02a397468970b5e06caa8de28fde40a Mon Sep 17 00:00:00 2001
From: Alexander J <741037+jaegeral@users.noreply.github.com>
Date: Tue, 10 Sep 2024 12:42:37 +0200
Subject: [PATCH] [CLI] export archive and unarchive a sketch (#3174)

* export archive and unarchive a sketch

* Add documentation to CLI

* feedback from review
---
 .../timesketch_cli_client/commands/sketch.py  | 65 +++++++++++++++++++
 docs/guides/user/cli-client.md                | 12 ++++
 2 files changed, 77 insertions(+)

diff --git a/cli_client/python/timesketch_cli_client/commands/sketch.py b/cli_client/python/timesketch_cli_client/commands/sketch.py
index 9850365199..11e0ee9671 100644
--- a/cli_client/python/timesketch_cli_client/commands/sketch.py
+++ b/cli_client/python/timesketch_cli_client/commands/sketch.py
@@ -13,11 +13,13 @@
 # limitations under the License.
 """Commands for sketches."""
 
+import time
 import json
 import click
 import pandas as pd
 
 from timesketch_cli_client.commands import attribute as attribute_command
+from timesketch_api_client import search
 
 
 @click.group("sketch")
@@ -93,3 +95,66 @@ def create_sketch(ctx, name, description):
         description = name
     sketch = api_client.create_sketch(name=name, description=description)
     click.echo(f"Sketch created: {sketch.name}")
+
+
+@sketch_group.command("export", help="Export a sketch")
+@click.option("--filename", required=True, help="Filename to export to.")
+@click.pass_context
+def export_sketch(ctx, filename):
+    """Export a sketch to a file.
+
+    Args:
+        filename: Filename to create
+    """
+    sketch = ctx.obj.sketch
+    click.echo("Executing export . . . ")
+    click.echo("Depending on the sketch size, this can take a while")
+    # start counting the time the export took
+    start_time = time.time()
+    try:
+        search_obj = search.Search(sketch=sketch)
+
+        click.echo(f"Number of events in that sketch: {search_obj.expected_size}")
+
+        search_obj.to_file(filename)
+        # Using the sketch.export function could be an alternative here
+        # TODO: https://github.com/google/timesketch/issues/2344
+        end_time = time.time()
+        click.echo(f"Export took {end_time - start_time} seconds")
+        click.echo("Finish")
+    except ValueError as e:
+        click.echo(f"Error: {e}")
+        ctx.exit(1)
+
+
+@sketch_group.command("archive", help="Archive a sketch")
+@click.pass_context
+def archive_sketch(ctx):
+    """Archive a sketch."""
+    sketch = ctx.obj.sketch
+    # if sketch is already archived error
+    if sketch.is_archived():
+        click.echo("Error Sketch is already archived")
+        ctx.exit(1)
+
+    # check if user has permissions
+    if not sketch.can_archive():
+        click.echo("User can not archive this sketch")
+        ctx.exit(1)
+
+    sketch.archive()
+    click.echo("Sketch archived")
+
+
+@sketch_group.command("unarchive", help="Unarchive a sketch")
+@click.pass_context
+def unarchive_sketch(ctx):
+    """Unarchive a sketch."""
+    sketch = ctx.obj.sketch
+    # if sketch is not archived error
+    if not sketch.is_archived():
+        click.echo("Error Sketch is not archived")
+        ctx.exit(1)
+    if sketch.is_archived():
+        sketch.unarchive()
+        click.echo("Sketch unarchived")
diff --git a/docs/guides/user/cli-client.md b/docs/guides/user/cli-client.md
index c7455a934a..c13a45c4f6 100644
--- a/docs/guides/user/cli-client.md
+++ b/docs/guides/user/cli-client.md
@@ -283,6 +283,18 @@ To remove an attribute from a sketch
 timesketch sketch remove_attribute
 ```
 
+### Archive Sketch
+
+Running `sketch archive` will set the archive flag to the sketch.
+
+### Unarchive a sketch
+
+Running `sketch unarchive` will set the archive flag to the sketch.
+
+### Export a sketch
+
+Running `sketch export` will export the complete Sketch to a file.
+
 ## Intelligence
 
 Intelligence is always sketch specific. The same can be achieved using