Fix CLI FileCategory choice breakage (#247)

* work around file category camelcase conversion to allow backwards compat values

* add camel-case Zip option

* make logic/tests work on current py42 as well as updated

* style

Author: timabrmsn

src/code42cli/click_ext/types.py

@@ -116,3 +116,19 @@ def _get_dt_from_date_time_pair(date, time):
             raise BadParameter("Unable to parse date string: {}.".format(date_string))
         else:
             return dt
+
+
+class MapChoice(click.Choice):
+    """Choice subclass that takes an extra map of additional 'valid' keys to map to correct
+    choices list, allowing backward compatible choice changes. The extra keys don't show up
+    in help text, but work when passed as a choice.
+    """
+
+    def __init__(self, choices, extras_map, **kwargs):
+        self.extras_map = extras_map
+        super().__init__(choices, **kwargs)
+
+    def convert(self, value, param, ctx):
+        if value in self.extras_map:
+            value = self.extras_map[value]
+        return super().convert(value, param, ctx)

src/code42cli/cmds/securitydata.py

@@ -13,6 +13,7 @@
 import code42cli.options as opt
 from code42cli.click_ext.groups import OrderedGroup
 from code42cli.click_ext.options import incompatible_with
+from code42cli.click_ext.types import MapChoice
 from code42cli.cmds.search import SendToCommand
 from code42cli.cmds.search.cursor_store import FileEventCursorStore
 from code42cli.cmds.search.extraction import handle_no_events
@@ -97,7 +98,25 @@
 file_category_option = click.option(
     "--file-category",
     multiple=True,
-    type=click.Choice(list(FileCategory.choices())),
+    type=MapChoice(
+        choices=list(FileCategory.choices()),
+        extras_map={
+            "AUDIO": FileCategory.AUDIO,
+            "DOCUMENT": FileCategory.DOCUMENT,
+            "EXECUTABLE": FileCategory.EXECUTABLE,
+            "IMAGE": FileCategory.IMAGE,
+            "PDF": FileCategory.PDF,
+            "PRESENTATION": FileCategory.PRESENTATION,
+            "SCRIPT": FileCategory.SCRIPT,
+            "SOURCE_CODE": FileCategory.SOURCE_CODE,
+            "SPREADSHEET": FileCategory.SPREADSHEET,
+            "VIDEO": FileCategory.VIDEO,
+            "VIRTUAL_DISK_IMAGE": FileCategory.VIRTUAL_DISK_IMAGE,
+            "ARCHIVE": FileCategory.ZIP,
+            "ZIP": FileCategory.ZIP,
+            "Zip": FileCategory.ZIP,
+        },
+    ),
     callback=searchopt.is_in_filter(f.FileCategory),
     cls=searchopt.AdvancedQueryAndSavedSearchIncompatible,
     help="Limits events to file events where the file can be classified by one of these categories.",

tests/cmds/test_securitydata.py

@@ -5,6 +5,7 @@
 import pytest
 from c42eventextractor.extractors import FileEventExtractor
 from py42.sdk.queries.fileevents.file_event_query import FileEventQuery
+from py42.sdk.queries.fileevents.filters.file_filter import FileCategory
 from tests.cmds.conftest import filter_term_is_in_call_args
 from tests.cmds.conftest import get_filter_value_from_json
 from tests.cmds.conftest import get_mark_for_search_and_send_to
@@ -650,7 +651,7 @@ def test_search_and_send_to_when_given_file_path_uses_file_path_filter(
 def test_search_and_send_to_when_given_file_category_uses_file_category_filter(
     runner, cli_state, file_event_extractor, command
 ):
-    file_category = "IMAGE"
+    file_category = FileCategory.IMAGE
     command = [*command, "--begin", "1h", "--file-category", file_category]
     runner.invoke(
         cli, command, obj=cli_state,
@@ -659,6 +660,44 @@ def test_search_and_send_to_when_given_file_category_uses_file_category_filter(
     assert str(f.FileCategory.is_in([file_category])) in filter_strings
 
 
+@pytest.mark.parametrize(
+    "category_choice",
+    [
+        ("AUDIO", FileCategory.AUDIO),
+        ("DOCUMENT", FileCategory.DOCUMENT),
+        ("EXECUTABLE", FileCategory.EXECUTABLE),
+        ("IMAGE", FileCategory.IMAGE),
+        ("PDF", FileCategory.PDF),
+        ("PRESENTATION", FileCategory.PRESENTATION),
+        ("SCRIPT", FileCategory.SCRIPT),
+        ("SOURCE_CODE", FileCategory.SOURCE_CODE),
+        ("SPREADSHEET", FileCategory.SPREADSHEET),
+        ("VIDEO", FileCategory.VIDEO),
+        ("VIRTUAL_DISK_IMAGE", FileCategory.VIRTUAL_DISK_IMAGE),
+        ("ARCHIVE", FileCategory.ZIP),
+        ("ZIP", FileCategory.ZIP),
+        ("Zip", FileCategory.ZIP),
+    ],
+)
+def test_all_caps_file_category_choices_convert_to_filecategory_constant(
+    runner, cli_state, file_event_extractor, category_choice
+):
+    ALL_CAPS_VALUE, camelCaseValue = category_choice
+    command = [
+        "security-data",
+        "search",
+        "--begin",
+        "1h",
+        "--file-category",
+        ALL_CAPS_VALUE,
+    ]
+    runner.invoke(
+        cli, command, obj=cli_state,
+    )
+    filter_strings = [str(arg) for arg in file_event_extractor.extract.call_args[0]]
+    assert str(f.FileCategory.is_in([camelCaseValue])) in filter_strings
+
+
 @search_and_send_to_test
 def test_search_and_send_to_when_given_process_owner_uses_process_owner_filter(
     runner, cli_state, file_event_extractor, command