Bugfix/search arg validation (#93)

* check against specific search_term enum instead of specified allowlist

* disallow begin/end/incremental

* remove unecessary changes

* remove unecessary changes

* remove unecessary changes

* get incompatible arg list from actual argument dict keys

* test and an h

* remove unused arg

Author: timabrmsn

src/code42cli/cmds/alerts/extraction.py

@@ -40,7 +40,7 @@ def extract(sdk, profile, output_logger, args):
     handlers = create_handlers(sdk, AlertExtractor, output_logger, store)
     extractor = AlertExtractor(sdk, handlers)
     if args.advanced_query:
-        exit_if_advanced_query_used_with_other_search_args(args)
+        exit_if_advanced_query_used_with_other_search_args(args, enums.AlertFilterArguments())
         extractor.extract_advanced(args.advanced_query)
     else:
         verify_begin_date_requirements(args, store)

src/code42cli/cmds/search_shared/args.py

@@ -3,7 +3,12 @@
 
 
 def create_search_args(search_for, filter_args):
-    search_args = {
+    advanced_query_incompatible_args = create_advanced_query_incompatible_search_args(search_for)
+    filter_args.update(advanced_query_incompatible_args)
+
+    format_enum = AlertOutputFormat() if search_for == "alerts" else OutputFormat()
+
+    advanced_query_compatible_args = {
         SearchArguments.ADVANCED_QUERY: ArgConfig(
             u"--{}".format(SearchArguments.ADVANCED_QUERY.replace(u"_", u"-")),
             metavar=u"QUERY_JSON",
@@ -13,6 +18,23 @@ def create_search_args(search_for, filter_args):
                 search_for
             ),
         ),
+        u"format": ArgConfig(
+            u"-f",
+            u"--format",
+            choices=format_enum,
+            default=format_enum.JSON,
+            help=u"The format used for outputting {0}.".format(search_for),
+        ),
+    }
+    filter_args.update(advanced_query_compatible_args)
+
+    return filter_args
+
+
+def create_advanced_query_incompatible_search_args(search_for=None):
+    """Returns a dict of args that are incompatible with the --advanced-query flag. Any new 
+    incompatible args should go here as this is function is also used for arg validation."""
+    args = {
         SearchArguments.BEGIN_DATE: ArgConfig(
             u"-b",
             u"--{}".format(SearchArguments.BEGIN_DATE),
@@ -30,23 +52,11 @@ def create_search_args(search_for, filter_args):
             help=u"The end of the date range in which to look for {0}, "
             u"argument format options are the same as --begin.".format(search_for),
         ),
-    }
-    format_enum = AlertOutputFormat() if search_for == "alerts" else OutputFormat()
-    format_and_incremental_args = {
-        u"format": ArgConfig(
-            u"-f",
-            u"--format",
-            choices=format_enum,
-            default=format_enum.JSON,
-            help=u"The format used for outputting {0}.".format(search_for),
-        ),
         u"incremental": ArgConfig(
             u"-i",
             u"--incremental",
             action=u"store_true",
             help=u"Only get {0} that were not previously retrieved.".format(search_for),
         ),
     }
-    search_args.update(filter_args)
-    search_args.update(format_and_incremental_args)
-    return search_args
+    return args

src/code42cli/cmds/search_shared/extraction.py

@@ -8,6 +8,7 @@
 from code42cli.logger import get_main_cli_logger
 from code42cli.cmds.alerts.util import get_alert_details
 from code42cli.util import warn_interrupt
+from code42cli.cmds.search_shared.args import create_advanced_query_incompatible_search_args
 
 logger = get_main_cli_logger()
 
@@ -73,13 +74,16 @@ def handle_response(response):
     return handlers
 
 
-def exit_if_advanced_query_used_with_other_search_args(args):
-    args_dict_copy = args.__dict__.copy()
-    for arg in (u"advanced_query", u"format", u"sdk", u"profile"):
-        args_dict_copy.pop(arg)
-    if any(args_dict_copy.values()):
-        logger.print_and_log_error(u"You cannot use --advanced-query with additional search args.")
-        exit(1)
+def exit_if_advanced_query_used_with_other_search_args(args, search_arg_enum):
+    incompatible_search_args_dict = create_advanced_query_incompatible_search_args()
+    incompatible_search_args_list = list(incompatible_search_args_dict.keys())
+    invalid_args = incompatible_search_args_list + list(search_arg_enum)
+    for arg in invalid_args:
+        if args.__dict__[arg]:
+            logger.print_and_log_error(
+                u"You cannot use --advanced-query with additional search args."
+            )
+            exit(1)
 
 
 def create_time_range_filter(filter_cls, begin_date=None, end_date=None):

src/code42cli/cmds/securitydata/extraction.py

@@ -1,7 +1,10 @@
 from c42eventextractor.extractors import FileEventExtractor
 from py42.sdk.queries.fileevents.filters import *
 
-from code42cli.cmds.search_shared.enums import ExposureType as ExposureTypeOptions
+from code42cli.cmds.search_shared.enums import (
+    ExposureType as ExposureTypeOptions,
+    FileEventFilterArguments,
+)
 from code42cli.cmds.search_shared.cursor_store import FileEventCursorStore
 from code42cli.cmds.search_shared.extraction import (
     verify_begin_date_requirements,
@@ -31,7 +34,7 @@ def extract(sdk, profile, output_logger, args):
     handlers = create_handlers(sdk, FileEventExtractor, output_logger, store)
     extractor = FileEventExtractor(sdk, handlers)
     if args.advanced_query:
-        exit_if_advanced_query_used_with_other_search_args(args)
+        exit_if_advanced_query_used_with_other_search_args(args, FileEventFilterArguments())
         extractor.extract_advanced(args.advanced_query)
     else:
         verify_begin_date_requirements(args, store)

tests/cmds/search_shared/test_advanced_query_args.py

@@ -0,0 +1,25 @@
+import pytest
+from code42cli.cmds.search_shared.extraction import (
+    exit_if_advanced_query_used_with_other_search_args,
+)
+from code42cli.cmds.search_shared.enums import FileEventFilterArguments, AlertFilterArguments
+
+
+def test_exit_if_advanced_query_provided_incompatible_args(
+    mocker, file_event_namespace, alert_namespace
+):
+    mock = mocker.patch(
+        "code42cli.cmds.search_shared.extraction.create_advanced_query_incompatible_search_args"
+    )
+    mock.return_value = {
+        "invalid_arg": None,
+    }
+    file_event_namespace.invalid_arg = "value"
+    with pytest.raises(SystemExit):
+        exit_if_advanced_query_used_with_other_search_args(
+            file_event_namespace, FileEventFilterArguments()
+        )
+
+    alert_namespace.invalid_arg = "value"
+    with pytest.raises(SystemExit):
+        exit_if_advanced_query_used_with_other_search_args(alert_namespace, AlertFilterArguments())