Armis collector enhancement (demisto#30491)

* initial commit * README changes * Added support for devices last run time, devices max fetch * Version bump and RN * Change unit test for new changes. Parse event type correctly. Make changes to comply with dict events instead of list events. * Changes to events from list to dict. * Remove redundant after AQL date filter * Alerts and Activity default max = 5k Devices default max = 10k * Remove redundant description in yml. * Updated the README * fix flake8 long line issues * Change param name to snake_case. No need to demisto.setLastRun multiple times. * Apply suggestions from code review Tech doc review changes. Co-authored-by: ShirleyDenkberg <[email protected]> * update RN. --------- Co-authored-by: ShirleyDenkberg <[email protected]>
code42 · Oct 31, 2023 · 5af4ebd · 5af4ebd
1 parent b540767
commit 5af4ebd
Show file tree

Hide file tree

Showing 6 changed files with 195 additions and 82 deletions.
diff --git a/Packs/Armis/Integrations/ArmisEventCollector/ArmisEventCollector.py b/Packs/Armis/Integrations/ArmisEventCollector/ArmisEventCollector.py
@@ -15,6 +15,7 @@ class EVENT_TYPE(NamedTuple):
     unique_id_key: str
     aql_query: str
     type: str
+    order_by: str
 
 
 ''' CONSTANTS '''
@@ -24,11 +25,14 @@ class EVENT_TYPE(NamedTuple):
 VENDOR = 'armis'
 PRODUCT = 'security'
 API_V1_ENDPOINT = '/api/v1'
-DEFAULT_MAX_FETCH = 1000
+DEFAULT_MAX_FETCH = 5000
+DEVICES_DEFAULT_MAX_FETCH = 10000
 EVENT_TYPES = {
-    'Alerts': EVENT_TYPE('alertId', 'in:alerts', 'alerts'),
-    'Threat activities': EVENT_TYPE('activityUUID', 'in:activity type:"Threat Detected"', 'threat_activities'),
+    'Alerts': EVENT_TYPE('alertId', 'in:alerts', 'alerts', 'time'),
+    'Activities': EVENT_TYPE('activityUUID', 'in:activity', 'activity', 'time'),
+    'Devices': EVENT_TYPE('id', 'in:devices', 'devices', 'firstSeen'),
 }
+DEVICES_LAST_FETCH = 'devices_last_fetch_time'
 
 ''' CLIENT CLASS '''
 
@@ -53,7 +57,8 @@ def update_access_token(self, access_token=None):
         self._headers = headers
         self._access_token = access_token
 
-    def fetch_by_aql_query(self, aql_query: str, max_fetch: int, after: None | datetime = None):
+    def fetch_by_aql_query(self, aql_query: str, max_fetch: int, after: None | datetime = None,
+                           order_by: str = 'time'):
         """ Fetches events using AQL query.
 
         Args:
@@ -64,7 +69,7 @@ def fetch_by_aql_query(self, aql_query: str, max_fetch: int, after: None | datet
         Returns:
             list[dict]: List of events objects represented as dictionaries.
         """
-        params: dict[str, Any] = {'aql': aql_query, 'includeTotal': 'true', 'length': max_fetch, 'orderBy': 'time'}
+        params: dict[str, Any] = {'aql': aql_query, 'includeTotal': 'true', 'length': max_fetch, 'orderBy': order_by}
         if not after:  # this should only happen when get-events command is used without from_date argument
             after = datetime.now()
         params['aql'] += f' after:{after.strftime(DATE_FORMAT)}'  # add 'after' date filter to AQL query in the desired format
@@ -264,7 +269,7 @@ def dedup_events(events: list[dict], events_last_fetch_ids: list[str], unique_id
         return new_events, new_ids
 
 
-def fetch_by_event_type(event_type: EVENT_TYPE, events: list, next_run: dict, client: Client,
+def fetch_by_event_type(event_type: EVENT_TYPE, events: dict, next_run: dict, client: Client,
                         max_fetch: int, last_run: dict, fetch_start_time: datetime | None):
     """ Fetch events by specific event type.
 
@@ -286,23 +291,28 @@ def fetch_by_event_type(event_type: EVENT_TYPE, events: list, next_run: dict, cl
     response = client.fetch_by_aql_query(
         aql_query=event_type.aql_query,
         max_fetch=max_fetch,
-        after=event_type_fetch_start_time
+        after=event_type_fetch_start_time,
+        order_by=event_type.order_by
     )
     demisto.debug(f'debug-log: fetched {len(response)} {event_type.type} from API')
     if response:
         new_events, next_run[last_fetch_ids] = dedup_events(
             response, last_run.get(last_fetch_ids, []), event_type.unique_id_key)
         next_run[last_fetch_time] = new_events[-1].get('time') if new_events else last_run.get(last_fetch_time)
-
-        events.extend(new_events)
+        events.setdefault(event_type.type, []).extend(new_events)
         demisto.debug(f'debug-log: overall {len(new_events)} {event_type.type} (after dedup)')
         demisto.debug(f'debug-log: last {event_type.type} in list: {new_events[-1] if new_events else {}}')
     else:
         next_run.update(last_run)
 
 
-def fetch_events(client: Client, max_fetch: int, last_run: dict, fetch_start_time: datetime | None,
-                 event_types_to_fetch: list[str]):
+def fetch_events(client: Client,
+                 max_fetch: int,
+                 devices_max_fetch: int,
+                 last_run: dict,
+                 fetch_start_time: datetime | None,
+                 event_types_to_fetch: list[str],
+                 device_fetch_interval: timedelta | None):
     """ Fetch events from Armis API.
 
     Args:
@@ -315,17 +325,22 @@ def fetch_events(client: Client, max_fetch: int, last_run: dict, fetch_start_tim
     Returns:
         (list[dict], dict) : List of fetched events and next run dictionary.
     """
-    events: list[dict] = []
+    events: dict[str, list[dict]] = {}
     next_run: dict[str, list | str] = {}
+    if 'Devices' in event_types_to_fetch\
+            and not should_run_device_fetch(last_run, device_fetch_interval, datetime.now()):
+        event_types_to_fetch.remove('Devices')
 
     for event_type in event_types_to_fetch:
+        event_max_fetch = max_fetch if event_type != "Devices" else devices_max_fetch
         try:
-            fetch_by_event_type(EVENT_TYPES[event_type], events, next_run, client, max_fetch, last_run, fetch_start_time)
+            fetch_by_event_type(EVENT_TYPES[event_type], events, next_run, client,
+                                event_max_fetch, last_run, fetch_start_time)
         except Exception as e:
             if "Invalid access token" in str(e):
                 client.update_access_token()
                 fetch_by_event_type(EVENT_TYPES[event_type], events, next_run, client,
-                                    max_fetch, last_run, fetch_start_time)
+                                    event_max_fetch, last_run, fetch_start_time)
 
     next_run['access_token'] = client._access_token
 
@@ -361,7 +376,8 @@ def handle_from_date_argument(from_date: str) -> datetime | None:
     return from_date_datetime if from_date_datetime else None
 
 
-def handle_fetched_events(events: list[dict[str, Any]], next_run: dict[str, str | list]):
+def handle_fetched_events(events: dict[str, list[dict[str, Any]]],
+                          next_run: dict[str, str | list]):
     """ Handle fetched events.
     - Send the fetched events to XSIAM.
     - Set last run values for next fetch cycle.
@@ -371,35 +387,41 @@ def handle_fetched_events(events: list[dict[str, Any]], next_run: dict[str, str
         next_run (dict[str, str | list]): Next run dictionary.
     """
     if events:
-        add_time_to_events(events)
-        demisto.debug(f'debug-log: {len(events)} events are about to be sent to XSIAM.')
-        send_events_to_xsiam(
-            events,
-            vendor=VENDOR,
-            product=PRODUCT
-        )
+        for event_type, events_list in events.items():
+            add_time_to_events(events_list)
+            demisto.debug(f'debug-log: {len(events_list)} events are about to be sent to XSIAM.')
+            product = f'{PRODUCT}_{event_type}' if event_type != 'alerts' else PRODUCT
+            send_events_to_xsiam(
+                events_list,
+                vendor=VENDOR,
+                product=product
+            )
+            demisto.debug(f'debug-log: {len(events)} events were sent to XSIAM.')
         demisto.setLastRun(next_run)
-        demisto.debug(f'debug-log: {len(events)} events were sent to XSIAM.')
     else:
         demisto.debug('debug-log: No new events fetched.')
 
     demisto.debug(f'debug-log: {next_run=}')
 
 
-def events_to_command_results(events: list[dict[str, Any]]) -> CommandResults:
+def events_to_command_results(events: dict[str, list[dict[str, Any]]]) -> list:
     """ Return a CommandResults object with a table of fetched events.
 
     Args:
         events (list[dict[str, Any]]): list of fetched events.
 
     Returns:
-        CommandResults: CommandResults object with a table of fetched events.
+        command_results_list: a List of CommandResults objects containing tables of fetched events.
     """
-    return CommandResults(
-        raw_response=events,
-        readable_output=tableToMarkdown(name=f'{VENDOR} {PRODUCT} events',
-                                        t=events,
-                                        removeNull=True))
+    command_results_list: list = []
+    for key, value in events.items():
+        product = f'{PRODUCT}_{key}' if key != 'alerts' else PRODUCT
+        command_results_list.append(CommandResults(
+            raw_response=events,
+            readable_output=tableToMarkdown(name=f'{VENDOR} {product} events',
+                                            t=value,
+                                            removeNull=True)))
+    return command_results_list
 
 
 def set_last_run_with_current_time(last_run: dict, event_types_to_fetch) -> None:
@@ -417,6 +439,29 @@ def set_last_run_with_current_time(last_run: dict, event_types_to_fetch) -> None
         last_run[last_fetch_time] = now_str
 
 
+def should_run_device_fetch(last_run,
+                            device_fetch_interval: timedelta | None,
+                            datetime_now: datetime):
+    """
+    Args:
+        last_run: last run object.
+        device_fetch_interval: device fetch interval.
+        datetime_now: time now
+
+    Returns: True if fetch device interval time has passed since last time that fetch run.
+
+    """
+    if not device_fetch_interval:
+        return False
+    if last_fetch_time := last_run.get(DEVICES_LAST_FETCH):
+        last_check_time = datetime.strptime(last_fetch_time, DATE_FORMAT)
+    else:
+        # first time device fetch
+        return True
+    demisto.debug(f'Should run device fetch? {last_check_time=}, {device_fetch_interval=}')
+    return datetime_now - last_check_time > device_fetch_interval
+
+
 ''' MAIN FUNCTION '''
 
 
@@ -430,11 +475,15 @@ def main():  # pragma: no cover
     base_url = urljoin(params.get('server_url'), API_V1_ENDPOINT)
     verify_certificate = not params.get('insecure', True)
     max_fetch = arg_to_number(params.get('max_fetch')) or DEFAULT_MAX_FETCH
+    devices_max_fetch = arg_to_number(params.get('devices_max_fetch')) or DEVICES_DEFAULT_MAX_FETCH
     proxy = params.get('proxy', False)
     event_types_to_fetch = argToList(params.get('event_types_to_fetch', []))
+    event_types_to_fetch = [event_type.strip(' ') for event_type in event_types_to_fetch]
     should_push_events = argToBoolean(args.get('should_push_events', False))
     from_date = args.get('from_date')
     fetch_start_time = handle_from_date_argument(from_date) if from_date else None
+    parsed_interval = dateparser.parse(params.get('deviceFetchInterval', '24 hours')) or dateparser.parse('24 hours')
+    device_fetch_interval: timedelta = (datetime.now() - parsed_interval)  # type: ignore[operator]
 
     demisto.debug(f'Command being called is {command}')
 
@@ -467,12 +516,14 @@ def main():  # pragma: no cover
                 events, next_run = fetch_events(
                     client=client,
                     max_fetch=max_fetch,
+                    devices_max_fetch=devices_max_fetch,
                     last_run=last_run,
                     fetch_start_time=fetch_start_time,
                     event_types_to_fetch=event_types_to_fetch,
+                    device_fetch_interval=device_fetch_interval
                 )
-
-                demisto.debug(f'debug-log: {len(events)} events fetched from armis api')
+                for key, value in events.items():
+                    demisto.debug(f'debug-log: {len(value)} events of type: {key} fetched from armis api')
 
                 if should_push_events:
                     handle_fetched_events(events, next_run)

diff --git a/Packs/Armis/Integrations/ArmisEventCollector/ArmisEventCollector.yml b/Packs/Armis/Integrations/ArmisEventCollector/ArmisEventCollector.yml
@@ -16,30 +16,53 @@ configuration:
   hiddenusername: true
   type: 9
   section: Connect
-- display: Number of events to fetch per type
+- display: Maximum number of events per fetch
   name: max_fetch
-  additionalinfo: The maximum number of events to fetch per event type.
+  additionalinfo: Alerts and activity events.
   type: 0
-  defaultvalue: 1000
+  defaultvalue: 5000
   section: Collect
+- display: Maximum number of device events per fetch
+  name: devices_max_fetch
+  type: 0
+  section: Collect
+  additionalinfo: Devices events.
+  defaultvalue: 10000
 - display: Trust any certificate (not secure)
   name: insecure
   type: 8
   section: Connect
 - display: Use system proxy settings
   name: proxy
-  type: 8
   section: Connect
+  type: 8
 - display: Event types to fetch
   name: event_types_to_fetch
   section: Collect
   required: true
   type: 16
-  defaultvalue: Alerts,Threat activities
+  defaultvalue: Alerts,Devices,Activities
   options:
   - Alerts
-  - Threat activities
-description: Collects alerts & threat activities from Armis resources.
+  - Devices
+  - Activities
+- section: Collect
+  advanced: true
+  display: Events Fetch Interval
+  additionalinfo: Alerts and activity events.
+  name: eventFetchInterval
+  defaultvalue: "1"
+  type: 19
+  required: false
+- section: Collect
+  advanced: true
+  display: Device Fetch Interval
+  additionalinfo: Time between fetch of devices (for example 12 hours, 60 minutes, etc.).
+  name: deviceFetchInterval
+  defaultvalue: "24 hours"
+  type: 0
+  required: false
+description: Collects alerts, devices and activities from Armis resources.
 display: Armis Event Collector
 name: ArmisEventCollector
 supportlevelheader: xsoar
@@ -54,7 +77,7 @@ script:
       - 'true'
       - 'false'
       required: true
-    - description: The date from which to fetch events. The format should be YYYY-MM-DD or YYYY-MM-DDT:HH:MM:SS. If not specified, the current date will be used.
+    - description: The date from which to fetch events. The format should be '20 minutes', '1 hour' or '2 days'.
       name: from_date
       required: false
     description: Manual command to fetch and display events. This command is used for developing/debugging and is to be used with caution, as it can create events, leading to events duplication and exceeding the API request limitation.