Merge pull request #9 from code42/add-results-param-to-get-all

Add results param to get all

Author: Juliya Smith

Packs/Code42/Integrations/Code42/Code42.py

@@ -187,12 +187,15 @@ def remove_user_from_departing_employee(self, username):
         self._get_sdk().detectionlists.departing_employee.remove(user_id)
         return user_id
 
-    def get_all_departing_employees(self):
+    def get_all_departing_employees(self, results):
         res = []
         pages = self._get_sdk().detectionlists.departing_employee.get_all()
         for page in pages:
             employees = page["items"]
-            res.extend(employees)
+            for employee in employees:
+                res.append(employee)
+                if len(res) == results:
+                    return res
         return res
 
     def add_user_to_high_risk_employee(self, username, note=None):
@@ -219,12 +222,16 @@ def remove_user_risk_tags(self, username, risk_tags):
         self._get_sdk().detectionlists.remove_user_risk_tags(user_id, risk_tags)
         return user_id
 
-    def get_all_high_risk_employees(self, risk_tags=None):
+    def get_all_high_risk_employees(self, risk_tags, results):
         risk_tags = _try_convert_str_list_to_list(risk_tags)
         res = []
         pages = self._get_sdk().detectionlists.high_risk_employee.get_all()
         for page in pages:
-            res.extend(_get_all_high_risk_employees_from_page(page, risk_tags))
+            employees = _get_all_high_risk_employees_from_page(page, risk_tags)
+            for employee in employees:
+                res.append(employee)
+                if len(res) == results:
+                    return res
         return res
 
     def fetch_alerts(self, start_time, event_severity_filter):
@@ -578,8 +585,9 @@ def departingemployee_remove_command(client, args):
 
 @logger
 def departingemployee_get_all_command(client, args):
+    results = args.get("results") or 50
     try:
-        employees = client.get_all_departing_employees()
+        employees = client.get_all_departing_employees(results)
         employees_context = [
             {
                 "UserID": e["userId"],
@@ -629,10 +637,11 @@ def highriskemployee_remove_command(client, args):
 @logger
 def highriskemployee_get_all_command(client, args):
     tags = args.get("risktags")
+    results = args.get("results")
     try:
-        employees = client.get_all_high_risk_employees(tags)
+        employees = client.get_all_high_risk_employees(tags, results)
         employees_context = [
-            {"UserID": e["userId"], "Username": e["userName"], "Note": e["notes"]}
+            {"UserID": e.get("userId"), "Username": e.get("userName"), "Note": e.get("notes")}
             for e in employees
         ]
         readable_outputs = tableToMarkdown("Retrieved All High Risk Employees", employees_context)
@@ -671,6 +680,36 @@ def highriskemployee_remove_risk_tags_command(client, args):
         return_error(create_command_error_message(demisto.command(), e))
 
 
+@logger
+def securitydata_search_command(client, args):
+    code42_security_data_context = []
+    _json = args.get("json")
+    file_context = []
+    # If JSON payload is passed as an argument, ignore all other args and search by JSON payload
+    if _json is not None:
+        file_events = client.search_file_events(_json)
+    else:
+        # Build payload
+        payload = build_query_payload(args)
+        file_events = client.search_file_events(payload)
+    if file_events:
+        for file_event in file_events:
+            code42_context_event = map_to_code42_event_context(file_event)
+            code42_security_data_context.append(code42_context_event)
+            file_context_event = map_to_file_context(file_event)
+            file_context.append(file_context_event)
+        readable_outputs = tableToMarkdown(
+            "Code42 Security Data Results",
+            code42_security_data_context,
+            headers=SECURITY_EVENT_HEADERS,
+        )
+        security_data_context_key = "Code42.SecurityData(val.EventID && val.EventID == obj.EventID)"
+        context = {security_data_context_key: code42_security_data_context, "File": file_context}
+        return readable_outputs, context, file_events
+    else:
+        return "No results found", {}, {}
+
+
 def _create_incident_from_alert_details(details):
     return {"name": "Code42 - {}".format(details["name"]), "occurred": details["createdAt"]}
 
@@ -680,7 +719,7 @@ def _stringify_lists_if_needed(event):
     shared_with = event.get("sharedWith")
     private_ip_addresses = event.get("privateIpAddresses")
     if shared_with:
-        shared_list = [u["cloudUsername"] for u in shared_with]
+        shared_list = [u.get("cloudUsername") for u in shared_with if u.get("cloudUsername")]
         event["sharedWith"] = str(shared_list)
     if private_ip_addresses:
         event["privateIpAddresses"] = str(private_ip_addresses)
@@ -729,7 +768,7 @@ def _fetch_remaining_incidents_from_last_run(self):
             if remaining_incidents:
                 return (
                     self._last_run,
-                    remaining_incidents[: self._fetch_limit],
+                    remaining_incidents[:self._fetch_limit],
                     remaining_incidents[self._fetch_limit:],
                 )
 
@@ -759,7 +798,11 @@ def _create_incident_from_alert(self, alert):
         return incident
 
     def _relate_files_to_alert(self, alert_details):
-        for obs in alert_details["observations"]:
+        observations = alert_details.get("observations")
+        if not observations:
+            alert_details["fileevents"] = []
+            return
+        for obs in observations:
             file_events = self._get_file_events_from_alert_details(obs, alert_details)
             alert_details["fileevents"] = [_process_event_from_observation(e) for e in file_events]
 
@@ -789,36 +832,6 @@ def fetch_incidents(
     return fetcher.fetch()
 
 
-@logger
-def securitydata_search_command(client, args):
-    code42_security_data_context = []
-    _json = args.get("json")
-    file_context = []
-    # If JSON payload is passed as an argument, ignore all other args and search by JSON payload
-    if _json is not None:
-        file_events = client.search_file_events(_json)
-    else:
-        # Build payload
-        payload = build_query_payload(args)
-        file_events = client.search_file_events(payload)
-    if file_events:
-        for file_event in file_events:
-            code42_context_event = map_to_code42_event_context(file_event)
-            code42_security_data_context.append(code42_context_event)
-            file_context_event = map_to_file_context(file_event)
-            file_context.append(file_context_event)
-        readable_outputs = tableToMarkdown(
-            "Code42 Security Data Results",
-            code42_security_data_context,
-            headers=SECURITY_EVENT_HEADERS,
-        )
-        security_data_context_key = "Code42.SecurityData(val.EventID && val.EventID == obj.EventID)"
-        context = {security_data_context_key: code42_security_data_context, "File": file_context}
-        return readable_outputs, context, file_events
-    else:
-        return "No results found", {}, {}
-
-
 def test_module(client):
     try:
         # Will fail if unauthorized

Packs/Code42/Integrations/Code42/Code42.yml

@@ -288,6 +288,11 @@ script:
       type: string
     description: Removes a user from the Departing Employee List.
   - name: code42-departingemployee-get-all
+    arguments:
+    - name: results
+      description: The number of items to return.
+      defaultvalue: "50"
+      type: number
     outputs:
     - contextPath: Code42.DepartingEmployee.UserID
       description: Internal Code42 User ID for the Departing Employee.
@@ -334,6 +339,10 @@ script:
     arguments:
     - name: risktags
       description: To filter results by employees who have these risk tags. Space delimited.
+    - name: results
+      description: The number of items to return.
+      defaultvalue: 50
+      type: number
     outputs:
     - contextPath: Code42.HighRiskEmployee.UserID
       description: Internal Code42 User ID for the High Risk Employee.

Packs/Code42/Integrations/Code42/Code42_test.py

@@ -1183,7 +1183,7 @@ def test_departingemployee_remove_command(code42_sdk_mock):
 
 def test_departingemployee_get_all_command(code42_departing_employee_mock):
     client = create_client(code42_departing_employee_mock)
-    _, _, res = departingemployee_get_all_command(client, {"username": "user1@example.com"})
+    _, _, res = departingemployee_get_all_command(client, {})
     expected = json.loads(MOCK_GET_ALL_DEPARTING_EMPLOYEES_RESPONSE)["items"]
     assert res == expected
     assert code42_departing_employee_mock.detectionlists.departing_employee.get_all.call_count == 1
@@ -1203,14 +1203,33 @@ def test_departingemployee_get_all_command_gets_employees_from_multiple_pages(
     )
     client = create_client(code42_departing_employee_mock)
 
-    _, _, res = departingemployee_get_all_command(client, {"username": "user1@example.com"})
+    _, _, res = departingemployee_get_all_command(client, {})
 
     # Expect to have employees from 3 pages in the result
     expected_page = json.loads(MOCK_GET_ALL_DEPARTING_EMPLOYEES_RESPONSE)["items"]
     expected = expected_page + expected_page + expected_page
     assert res == expected
 
 
+def test_departingemployee_get_all_command_gets_number_of_employees_equal_to_results_param(
+    code42_departing_employee_mock, mocker
+):
+
+    # Setup get all departing employees
+    page = MOCK_GET_ALL_DEPARTING_EMPLOYEES_RESPONSE
+    # Setup 3 pages of employees
+    employee_page_generator = (
+        create_mock_code42_sdk_response(mocker, page) for page in [page, page, page]
+    )
+    code42_departing_employee_mock.detectionlists.departing_employee.get_all.return_value = (
+        employee_page_generator
+    )
+    client = create_client(code42_departing_employee_mock)
+
+    _, _, res = departingemployee_get_all_command(client, {"results": 1})
+    assert len(res) == 1
+
+
 def test_departingemployee_get_all_command_when_no_employees(
     code42_departing_employee_mock, mocker
 ):
@@ -1260,29 +1279,6 @@ def test_highriskemployee_remove_command(code42_sdk_mock):
     code42_sdk_mock.detectionlists.high_risk_employee.remove.assert_called_once_with(expected)
 
 
-def test_fetch_when_no_significant_file_categories_ignores_filter(
-    code42_fetch_incidents_mock, mocker
-):
-    response_text = MOCK_ALERT_DETAILS_RESPONSE.replace(
-        '"isSignificant": true', '"isSignificant": false'
-    )
-    alert_details_response = create_mock_code42_sdk_response(mocker, response_text)
-    code42_fetch_incidents_mock.alerts.get_details.return_value = alert_details_response
-    client = create_client(code42_fetch_incidents_mock)
-    _, _, _ = fetch_incidents(
-        client=client,
-        last_run={"last_fetch": None},
-        first_fetch_time=MOCK_FETCH_TIME,
-        event_severity_filter=None,
-        fetch_limit=10,
-        include_files=True,
-        integration_context=None,
-    )
-    actual_query = str(code42_fetch_incidents_mock.securitydata.search_file_events.call_args[0][0])
-    assert "fileCategory" not in actual_query
-    assert "IMAGE" not in actual_query
-
-
 def test_highriskemployee_get_all_command(code42_high_risk_employee_mock):
     client = create_client(code42_high_risk_employee_mock)
     _, _, res = highriskemployee_get_all_command(client, {})
@@ -1327,6 +1323,23 @@ def test_highriskemployee_get_all_command_when_given_risk_tags_only_gets_employe
     assert code42_high_risk_employee_mock.detectionlists.high_risk_employee.get_all.call_count == 1
 
 
+def test_highriskemployee_get_all_command_gets_number_of_employees_equal_to_results_param(
+    code42_high_risk_employee_mock, mocker
+):
+    # Setup get all high risk employees
+    page = MOCK_GET_ALL_HIGH_RISK_EMPLOYEES_RESPONSE
+    # Setup 3 pages of employees
+    employee_page_generator = (
+        create_mock_code42_sdk_response(mocker, page) for page in [page, page, page]
+    )
+    code42_high_risk_employee_mock.detectionlists.high_risk_employee.get_all.return_value = (
+        employee_page_generator
+    )
+    client = create_client(code42_high_risk_employee_mock)
+    _, _, res = highriskemployee_get_all_command(client, {"results": 1})
+    assert len(res) == 1
+
+
 def test_highriskemployee_get_all_command_when_no_employees(code42_high_risk_employee_mock, mocker):
     no_employees_response = get_empty_detectionlist_response(
         mocker, MOCK_GET_ALL_HIGH_RISK_EMPLOYEES_RESPONSE
@@ -1387,6 +1400,29 @@ def test_security_data_search_command(code42_file_events_mock):
     assert filter_groups[3]["filters"][0]["value"] == "ApplicationRead"
 
 
+def test_fetch_when_no_significant_file_categories_ignores_filter(
+    code42_fetch_incidents_mock, mocker
+):
+    response_text = MOCK_ALERT_DETAILS_RESPONSE.replace(
+        '"isSignificant": true', '"isSignificant": false'
+    )
+    alert_details_response = create_mock_code42_sdk_response(mocker, response_text)
+    code42_fetch_incidents_mock.alerts.get_details.return_value = alert_details_response
+    client = create_client(code42_fetch_incidents_mock)
+    _, _, _ = fetch_incidents(
+        client=client,
+        last_run={"last_fetch": None},
+        first_fetch_time=MOCK_FETCH_TIME,
+        event_severity_filter=None,
+        fetch_limit=10,
+        include_files=True,
+        integration_context=None,
+    )
+    actual_query = str(code42_fetch_incidents_mock.securitydata.search_file_events.call_args[0][0])
+    assert "fileCategory" not in actual_query
+    assert "IMAGE" not in actual_query
+
+
 def test_fetch_incidents_handles_single_severity(code42_sdk_mock):
     client = create_client(code42_sdk_mock)
     fetch_incidents(