Add bulk counterparts for add-risk-tags and remove-risk-tags (#77)

* Added bulk feature in hre to add or remove risk-tags

* Add bulk generate template for high risk employee

* Added tests and rectify inheritance syntax

* Added changelog

* Docs

* Refactor

* add extra space

* Update CL

Co-authored-by: Juliya Smith <juliya.smith@code42.com>

Author: kiran-chaudhary

CHANGELOG.md

@@ -20,6 +20,11 @@ how a consumer would use the library (e.g. adding unit tests, updating documenta
 
 ### Added
 
+- `code42 high-risk-employee bulk` supports `add-risk-tags` and `remove-risk-tags`.
+    - `code42 high-risk-employee bulk generate-template <cmd>` options `add-risk-tags` and `remove-risk-tags`.
+        - `add-risk-tags` that takes a csv file with username and space separated risk tags.
+        - `remove-risk-tags` that takes a csv file with username and space separated risk tags.
+
 - Display, `Fuzzy suggestions`, valid keywords matching mistyped commands or arguments.
 
 - `code42 alerts`:

src/code42cli/cmds/detectionlists/__init__.py

@@ -1,10 +1,11 @@
 from py42.exceptions import Py42BadRequestError
 
-from code42cli.bulk import generate_template, run_bulk_process, BulkCommandType
+from code42cli.bulk import generate_template, run_bulk_process
 from code42cli.file_readers import create_csv_reader, create_flat_file_reader
 from code42cli.errors import UserAlreadyAddedError, UserDoesNotExistError, UnknownRiskTagError
-from code42cli.cmds.detectionlists.commands import DetectionListCommandFactory
 from code42cli.cmds.detectionlists.enums import DetectionLists, DetectionListUserKeys, RiskTags
+from code42cli.cmds.detectionlists.commands import DetectionListCommandFactory
+from code42cli.cmds.detectionlists.bulk import BulkDetectionList, BulkHighRiskEmployee
 
 
 def try_handle_user_already_added_error(bad_request_err, username_tried_adding, list_name):
@@ -31,6 +32,9 @@ def __init__(self, add=None, remove=None, load_add=None):
         self.remove_employee = remove
         self.load_add_description = load_add
 
+    def add_handler(self, attr_name, handler):
+        self.__setattr__(attr_name, handler)
+
 
 class DetectionList(object):
     """An object representing a Code42 detection list. Use this class by passing in handlers for 
@@ -88,29 +92,49 @@ def load_subcommands(self):
         return [bulk, add, remove]
 
     def _load_bulk_subcommands(self):
-        generate_template_cmd = self.factory.create_bulk_generate_template_command(
-            self.generate_template_file
-        )
+
         add = self.factory.create_bulk_add_command(self.bulk_add_employees)
         remove = self.factory.create_bulk_remove_command(self.bulk_remove_employees)
-        return [generate_template_cmd, add, remove]
+        commands = [add, remove]
+
+        if self.name == DetectionLists.HIGH_RISK_EMPLOYEE:
+            commands.extend(self._get_risk_tags_bulk_subcommands())
+        else:
+            generate_template_cmd = self.factory.create_bulk_generate_template_command(
+                self.generate_template_file
+            )
+            commands.append(generate_template_cmd)
+        return commands
+
+    def _get_risk_tags_bulk_subcommands(self):
+        bulk_add_risk_tags = self.factory.create_bulk_add_risk_tags_command(self.bulk_add_risk_tags)
+        bulk_remove_risk_tags = self.factory.create_bulk_remove_risk_tags_command(
+            self.bulk_remove_risk_tags
+        )
+
+        self.handlers.add_handler(u"add_risk_tags", add_risk_tags)
+        self.handlers.add_handler(u"remove_risk_tags", remove_risk_tags)
+        generate_template_cmd = self.factory.create_hre_bulk_generate_template_command(
+            self.generate_template_file
+        )
+        return [bulk_add_risk_tags, bulk_remove_risk_tags, generate_template_cmd]
 
     def generate_template_file(self, cmd, path=None):
         """Generates a template file a user would need to fill-in for bulk operating on the 
         detection list.
-        
+
         Args:
             cmd (str or unicode): An option from the `BulkCommandType` enum specifying which type of file to 
                 generate.
             path (str or unicode, optional): A path to put the file after it's generated. If None, will use 
                 the current working directory. Defaults to None.
         """
-        handler = None
-        if cmd == BulkCommandType.ADD:
-            handler = self.handlers.add_employee
-        elif cmd == BulkCommandType.REMOVE:
-            handler = self.handlers.remove_employee
 
+        if self.name == DetectionLists.HIGH_RISK_EMPLOYEE:
+            detection_list = BulkHighRiskEmployee()
+        else:
+            detection_list = BulkDetectionList()
+        handler = detection_list.get_handler(self.handlers, cmd)
         generate_template(handler, path)
 
     def bulk_add_employees(self, sdk, profile, csv_file):
@@ -145,6 +169,14 @@ def _add_employee(self, sdk, profile, **kwargs):
     def _remove_employee(self, sdk, profile, *args, **kwargs):
         self.handlers.remove_employee(sdk, profile, *args, **kwargs)
 
+    def bulk_add_risk_tags(self, sdk, profile, csv_file):
+        reader = create_csv_reader(csv_file)
+        run_bulk_process(lambda **kwargs: add_risk_tags(sdk, profile, **kwargs), reader)
+
+    def bulk_remove_risk_tags(self, sdk, profile, csv_file):
+        reader = create_csv_reader(csv_file)
+        run_bulk_process(lambda **kwargs: remove_risk_tags(sdk, profile, **kwargs), reader)
+
 
 def load_username_description(argument_collection):
     """Loads the arg descriptions for the `username` CLI parameter."""
@@ -224,3 +256,23 @@ def _try_handle_bad_risk_tag(tags):
     unknowns = [tag for tag in tags if tag not in options] if tags else None
     if unknowns:
         raise UnknownRiskTagError(unknowns)
+
+
+def handle_list_args(list_arg):
+    """Converts str args to a list. Useful for `bulk` commands which don't use `argparse` but
+    instead pass in values from files, such as in the form "item1 item2"."""
+    if list_arg and not isinstance(list_arg, list):
+        return list_arg.split()
+    return list_arg
+
+
+def add_risk_tags(sdk, profile, username, tag):
+    risk_tag = handle_list_args(tag)
+    user_id = get_user_id(sdk, username)
+    try_add_risk_tags(sdk, user_id, risk_tag)
+
+
+def remove_risk_tags(sdk, profile, username, tag):
+    risk_tag = handle_list_args(tag)
+    user_id = get_user_id(sdk, username)
+    try_remove_risk_tags(sdk, user_id, risk_tag)

src/code42cli/cmds/detectionlists/bulk.py

@@ -0,0 +1,41 @@
+from code42cli.bulk import BulkCommandType
+
+
+class HighRiskBulkCommandType(BulkCommandType):
+    ADD_RISK_TAG = u"add-risk-tags"
+    REMOVE_RISK_TAG = u"remove-risk-tags"
+
+    def __iter__(self):
+        parent_items = list(super(HighRiskBulkCommandType, self).__iter__())
+        return iter([parent_items[0], parent_items[1], self.ADD_RISK_TAG, self.REMOVE_RISK_TAG])
+
+
+class BulkDetectionList(object):
+
+    def __init__(self):
+        self.type = BulkCommandType
+        
+    def get_handler(self, handlers, cmd):
+        handler = None
+        if cmd == self.type.ADD:
+            handler = handlers.add_employee
+        elif cmd == self.type.REMOVE:
+            handler = handlers.remove_employee
+        return handler
+
+
+class BulkHighRiskEmployee(BulkDetectionList):
+
+    def __init__(self):
+        super(BulkHighRiskEmployee, self).__init__()
+        self.type = HighRiskBulkCommandType
+
+    def get_handler(self, handlers, cmd):
+        handler = super(BulkHighRiskEmployee, self).get_handler(handlers, cmd)
+        if not handler:
+            if cmd == self.type.ADD_RISK_TAG:
+                handler = handlers.add_risk_tags
+            elif cmd == self.type.REMOVE_RISK_TAG:
+                handler = handlers.remove_risk_tags
+
+        return handler

src/code42cli/cmds/detectionlists/commands.py

@@ -1,5 +1,6 @@
 from code42cli.bulk import BulkCommandType
 from code42cli.commands import Command
+from code42cli.cmds.detectionlists.bulk import HighRiskBulkCommandType
 
 
 def create_usage_prefix(detection_list_name):
@@ -52,6 +53,15 @@ def create_bulk_generate_template_command(self, handler):
             arg_customizer=DetectionListCommandFactory._load_bulk_generate_template_description,
         )
 
+    def create_hre_bulk_generate_template_command(self, handler):
+        return Command(
+            u"generate-template",
+            u"Generate the necessary csv template needed for bulk adding users.",
+            u"{} generate-template <cmd> <optional args>".format(self._bulk_usage_prefix),
+            handler=handler,
+            arg_customizer=DetectionListCommandFactory._load_hre_bulk_generate_template_description,
+        )
+
     def create_bulk_add_command(self, handler):
         return Command(
             BulkCommandType.ADD,
@@ -70,12 +80,36 @@ def create_bulk_remove_command(self, handler):
             arg_customizer=self._load_bulk_remove_description,
         )
 
+    def create_bulk_add_risk_tags_command(self, handler):
+        return Command(
+            u"add-risk-tags",
+            u"Associates risk tags with a user in bulk.",
+            u"{} {} <file>".format(self._bulk_usage_prefix, HighRiskBulkCommandType.ADD_RISK_TAG),
+            handler=handler,
+            arg_customizer=self._load_bulk_add_risk_tags_description,
+        )
+
+    def create_bulk_remove_risk_tags_command(self, handler):
+        return Command(
+            u"remove-risk-tags",
+            u"Disassociates risk tags from a user in bulk.",
+            u"{} {} <file>".format(self._bulk_usage_prefix, HighRiskBulkCommandType.REMOVE_RISK_TAG),
+            handler=handler,
+            arg_customizer=self._load_bulk_remove_risk_tags_description,
+        )
+
     @staticmethod
     def _load_bulk_generate_template_description(argument_collection):
         cmd_type = argument_collection.arg_configs[u"cmd"]
-        cmd_type.set_help(u"The type of command the template with be used for.")
+        cmd_type.set_help(u"The type of command the template will be used for.")
         cmd_type.set_choices(BulkCommandType())
 
+    @staticmethod
+    def _load_hre_bulk_generate_template_description(argument_collection):
+        cmd_type = argument_collection.arg_configs[u"cmd"]
+        cmd_type.set_help(u"The type of command the template will be used for.")
+        cmd_type.set_choices(HighRiskBulkCommandType())
+
     def _load_bulk_add_description(self, argument_collection):
         csv_file = argument_collection.arg_configs[u"csv_file"]
         csv_file.set_help(
@@ -91,3 +125,23 @@ def _load_bulk_remove_description(self, argument_collection):
                 self._name
             )
         )
+
+    def _load_bulk_add_risk_tags_description(self, argument_collection):
+        csv_file = argument_collection.arg_configs[u"csv_file"]
+        csv_file.set_help(
+            u"A file containing a ',' separated username with space-separated tags to add "
+            u"to the {} detection list. "
+            u"e.g. test@email.com,tag1 tag2 tag3".format(
+                self._name
+            )
+        )
+
+    def _load_bulk_remove_risk_tags_description(self, argument_collection):
+        csv_file = argument_collection.arg_configs[u"csv_file"]
+        csv_file.set_help(
+            u"A file containing a ',' separated username with space-separated tags to remove "
+            u"from the {} detection list. "
+            u"e.g. test@email.com,tag1 tag2 tag3".format(
+                self._name
+            )
+        )

src/code42cli/cmds/detectionlists/high_risk_employee.py

@@ -1,23 +1,25 @@
+from py42.exceptions import Py42BadRequestError
+
+from code42cli.commands import Command
 from code42cli.cmds.detectionlists import (
     DetectionList,
     DetectionListHandlers,
     load_user_descriptions,
-    load_username_description,
     get_user_id,
     update_user,
     try_handle_user_already_added_error,
-    try_add_risk_tags,
-    try_remove_risk_tags,
+    add_risk_tags,
+    remove_risk_tags,
+    load_username_description,
+    handle_list_args,
 )
 from code42cli.cmds.detectionlists.enums import DetectionLists, DetectionListUserKeys, RiskTags
-from code42cli.commands import Command
-
-from py42.exceptions import Py42BadRequestError
 
 
 def load_subcommands():
     handlers = _create_handlers()
     detection_list = DetectionList.create_high_risk_employee_list(handlers)
+
     cmd_list = detection_list.load_subcommands()
     cmd_list.extend(
         [
@@ -26,14 +28,14 @@ def load_subcommands():
                 u"Associates risk tags with a user.",
                 u"code42 high-risk-employee add-risk-tags --username <username> --tag <risk-tags>",
                 handler=add_risk_tags,
-                arg_customizer=_load_risk_tag_mgmt_descriptions,
+                arg_customizer=load_risk_tag_mgmt_descriptions,
             ),
             Command(
                 u"remove-risk-tags",
                 u"Disassociates risk tags from a user.",
                 u"code42 high-risk-employee remove-risk-tags --username <username> --tag <risk-tags>",
                 handler=remove_risk_tags,
-                arg_customizer=_load_risk_tag_mgmt_descriptions,
+                arg_customizer=load_risk_tag_mgmt_descriptions,
             ),
         ]
     )
@@ -46,18 +48,6 @@ def _create_handlers():
     )
 
 
-def add_risk_tags(sdk, profile, username, tag):
-    risk_tag = _handle_list_args(tag)
-    user_id = get_user_id(sdk, username)
-    try_add_risk_tags(sdk, user_id, risk_tag)
-
-
-def remove_risk_tags(sdk, profile, username, tag):
-    risk_tag = _handle_list_args(tag)
-    user_id = get_user_id(sdk, username)
-    try_remove_risk_tags(sdk, user_id, risk_tag)
-
-
 def add_high_risk_employee(sdk, profile, username, cloud_alias=None, risk_tag=None, notes=None):
     """Adds an employee to the high risk employee detection list.
     
@@ -69,7 +59,7 @@ def add_high_risk_employee(sdk, profile, username, cloud_alias=None, risk_tag=No
         risk_tag (iter[str]): Risk tags associated with the employee.
         notes: (str): Notes about the employee.
     """
-    risk_tag = _handle_list_args(risk_tag)
+    risk_tag = handle_list_args(risk_tag)
     user_id = get_user_id(sdk, username)
 
     try:
@@ -93,7 +83,12 @@ def remove_high_risk_employee(sdk, profile, username):
     sdk.detectionlists.high_risk_employee.remove(user_id)
 
 
-def _load_risk_tag_description(argument_collection):
+def _load_add_description(argument_collection):
+    load_user_descriptions(argument_collection)
+    load_risk_tag_description(argument_collection)
+
+
+def load_risk_tag_description(argument_collection):
     risk_tag = (
         argument_collection.arg_configs.get(DetectionListUserKeys.RISK_TAG)
         or argument_collection.arg_configs[u"tag"]
@@ -105,19 +100,6 @@ def _load_risk_tag_description(argument_collection):
     )
 
 
-def _load_add_description(argument_collection):
-    load_user_descriptions(argument_collection)
-    _load_risk_tag_description(argument_collection)
-
-
-def _load_risk_tag_mgmt_descriptions(argument_collection):
+def load_risk_tag_mgmt_descriptions(argument_collection):
     load_username_description(argument_collection)
-    _load_risk_tag_description(argument_collection)
-
-
-def _handle_list_args(list_arg):
-    """Converts str args to a list. Useful for `bulk` commands which don't use `argparse` but 
-    instead pass in values from files, such as in the form "item1 item2"."""
-    if list_arg and type(list_arg) != list:
-        return list_arg.split()
-    return list_arg
+    load_risk_tag_description(argument_collection)