DHS Feed v2 (#19760)

* add v2

Packs/ApiModules/ReleaseNotes/2_2_12.md

@@ -0,0 +1,5 @@
+
+#### Scripts
+##### TAXII2ApiModule
+- Added support for limiting the number of fetched indicators.
+- Improved implementation for polling collections.

Packs/ApiModules/Scripts/TAXII2ApiModule/TAXII2ApiModule.py

@@ -2,7 +2,7 @@
 from CommonServerPython import *
 from CommonServerUserPython import *
 
-from typing import Union, Optional, List, Dict, Tuple
+from typing import Optional, List, Dict, Tuple
 from requests.sessions import merge_setting, CaseInsensitiveDict
 import re
 import copy
@@ -25,8 +25,6 @@
 
 ERR_NO_COLL = "No collection is available for this user, please make sure you entered the configuration correctly"
 
-DATE_FORMAT = '%Y-%m-%dT%H:%M:%S.%fZ'
-
 # Pattern Regexes - used to extract indicator type and value
 INDICATOR_OPERATOR_VAL_FORMAT_PATTERN = r"(\w.*?{value}{operator})'(.*?)'"
 INDICATOR_EQUALS_VAL_PATTERN = INDICATOR_OPERATOR_VAL_FORMAT_PATTERN.format(
@@ -175,6 +173,10 @@
                           'ZA': 'South Africa', 'ZM': 'Zambia', 'ZW': 'Zimbabwe'}
 
 
+def reached_limit(limit: int, element_count: int):
+    return element_count >= limit > -1
+
+
 class Taxii2FeedClient:
     def __init__(
             self,
@@ -314,12 +316,18 @@ def init_roots(self):
             logging.disable(logging.NOTSET)
 
     def set_api_root(self):
-        roots_to_api = {str(api_root.url).split('/')[-2]: api_root
-                        for api_root in self.server.api_roots}  # type: ignore[attr-defined]
+        roots_to_api = {}
+        for api_root in self.server.api_roots:  # type: ignore[attr-defined]
+            # ApiRoots are initialized with wrong _conn because we are not providing auth or cert to Server
+            # closing wrong unused connections
+            api_root_name = str(api_root.url).split('/')[-2]
+            demisto.debug(f'closing api_root._conn for {api_root_name}')
+            api_root._conn.close()
+            roots_to_api[api_root_name] = api_root
 
         if self.default_api_root:
             if not roots_to_api.get(self.default_api_root):
-                raise DemistoException(f'The given default API root {self.default_api_root} doesn\'t exists.'
+                raise DemistoException(f'The given default API root {self.default_api_root} doesn\'t exist. '
                                        f'Available API roots are {list(roots_to_api.keys())}.')
             self.api_root = roots_to_api.get(self.default_api_root)
 
@@ -1001,7 +1009,6 @@ def build_iterator(self, limit: int = -1, **kwargs) -> List[Dict[str, str]]:
         :param limit: max amount of indicators to fetch
         :return: Cortex indicators list
         """
-
         if not isinstance(self.collection_to_fetch, (v20.Collection, v21.Collection)):
             raise DemistoException(
                 "Could not find a collection to fetch from. "
@@ -1018,7 +1025,7 @@ def build_iterator(self, limit: int = -1, **kwargs) -> List[Dict[str, str]]:
 
         return indicators
 
-    def load_stix_objects_from_envelope(self, envelopes: Dict[str, Any], limit: int = -1):
+    def load_stix_objects_from_envelope(self, envelopes: types.GeneratorType, limit: int = -1):
 
         parse_stix_2_objects = {
             "indicator": self.parse_indicator,
@@ -1045,118 +1052,65 @@ def load_stix_objects_from_envelope(self, envelopes: Dict[str, Any], limit: int
             "location": self.parse_location,
             "vulnerability": self.parse_vulnerability
         }
-        indicators = []
 
-        # TAXII 2.0
-        if isinstance(list(envelopes.values())[0], types.GeneratorType):
-            indicators.extend(self.parse_generator_type_envelope(envelopes, parse_stix_2_objects))
-        # TAXII 2.1
-        else:
-            indicators.extend(self.parse_dict_envelope(envelopes, parse_stix_2_objects, limit))
+        indicators, relationships_lst = self.parse_generator_type_envelope(envelopes, parse_stix_2_objects, limit)
+        if relationships_lst:
+            indicators.extend(self.parse_relationships(relationships_lst))
         demisto.debug(
             f"TAXII 2 Feed has extracted {len(indicators)} indicators"
         )
-        if limit > -1:
-            return indicators[:limit]
+
         return indicators
 
-    def parse_generator_type_envelope(self, envelopes: Dict[str, Any],
-                                      parse_objects_func):
+    def parse_generator_type_envelope(self, envelopes: types.GeneratorType, parse_objects_func, limit: int = -1):
         indicators = []
         relationships_lst = []
-        for obj_type, envelope in envelopes.items():
-            for sub_envelope in envelope:
-                stix_objects = sub_envelope.get("objects")
-                if not stix_objects:
-                    # no fetched objects
-                    break
-                # now we have a list of objects, go over each obj, save id with obj, parse the obj
-                if obj_type != "relationship":
-                    for obj in stix_objects:
-                        # we currently don't support extension object
-                        if obj.get('type') == 'extension-definition':
-                            continue
-                        self.id_to_object[obj.get('id')] = obj
-                        result = parse_objects_func[obj_type](obj)
-                        if not result:
-                            continue
-                        indicators.extend(result)
-                        self.update_last_modified_indicator_date(obj.get("modified"))
-                else:
-                    relationships_lst.extend(stix_objects)
-        if relationships_lst:
-            indicators.extend(self.parse_relationships(relationships_lst))
+        for envelope in envelopes:
+            stix_objects = envelope.get("objects")
+            if not stix_objects:
+                # no fetched objects
+                break
 
-        return indicators
+            # now we have a list of objects, go over each obj, save id with obj, parse the obj
+            for obj in stix_objects:
+                obj_type = obj.get('type')
+
+                # we currently don't support extension object
+                if obj_type == 'extension-definition':
+                    continue
+                elif obj_type == 'relationship':
+                    relationships_lst.append(obj)
+                    continue
 
-    def parse_dict_envelope(self, envelopes: Dict[str, Any],
-                            parse_objects_func, limit: int = -1):
-        indicators: list = []
-        relationships_list: List[Dict[str, Any]] = []
-        for obj_type, envelope in envelopes.items():
-            cur_limit = limit
-            stix_objects = envelope.get("objects", [])
-            if obj_type != "relationship":
-                for obj in stix_objects:
-                    # we currently don't support extension object
-                    if obj.get('type') == 'extension-definition':
-                        continue
-                    self.id_to_object[obj.get('id')] = obj
-                    result = parse_objects_func[obj_type](obj)
-                    if not result:
-                        continue
+                self.id_to_object[obj.get('id')] = obj
+                if not parse_objects_func.get(obj_type):
+                    demisto.debug(f'There is no parsing function for object type {obj_type}, '
+                                  f'available parsing functions are for types: {",".join(parse_objects_func.keys())}.')
+                    continue
+                if result := parse_objects_func[obj_type](obj):
                     indicators.extend(result)
                     self.update_last_modified_indicator_date(obj.get("modified"))
-            else:
-                relationships_list.extend(stix_objects)
 
-            while envelope.get("more", False):
-                page_size = self.get_page_size(limit, cur_limit)
-                envelope = self.collection_to_fetch.get_objects(
-                    limit=page_size, next=envelope.get("next", ""), type=obj_type
-                )
-                if isinstance(envelope, Dict):
-                    stix_objects = envelope.get("objects")
-                    if obj_type != "relationship":
-                        for obj in stix_objects:
-                            self.id_to_object[obj.get('id')] = obj
-                            result = parse_objects_func[obj_type](obj)
-                            if not result:
-                                continue
-                            indicators.extend(result)
-                            self.update_last_modified_indicator_date(obj.get("modified"))
-                    else:
-                        relationships_list.extend(stix_objects)
-                else:
-                    raise DemistoException(
-                        "Error: TAXII 2 client received the following response while requesting "
-                        f"indicators: {str(envelope)}\n\nExpected output is json"
-                    )
+                if reached_limit(limit, len(indicators)):
+                    return indicators, relationships_lst
 
-        if relationships_list:
-            indicators.extend(self.parse_relationships(relationships_list))
-        return indicators
+        return indicators, relationships_lst
 
     def poll_collection(
             self, page_size: int, **kwargs
-    ) -> Dict[str, Union[types.GeneratorType, Dict[str, str]]]:
+    ) -> types.GeneratorType:
         """
         Polls a taxii collection
         :param page_size: size of the request page
         """
-        types_envelopes = {}
         get_objects = self.collection_to_fetch.get_objects
-        if len(self.objects_to_fetch) > 1:  # when fetching one type no need to fetch relationship
+        if 'relationship' not in self.objects_to_fetch and \
+                len(self.objects_to_fetch) > 1:  # when fetching one type no need to fetch relationship
             self.objects_to_fetch.append('relationship')
-        for obj_type in self.objects_to_fetch:
-            kwargs['type'] = obj_type
-            if isinstance(self.collection_to_fetch, v20.Collection):
-                envelope = v20.as_pages(get_objects, per_request=page_size, **kwargs)
-            else:
-                envelope = get_objects(limit=page_size, **kwargs)
-            if envelope:
-                types_envelopes[obj_type] = envelope
-        return types_envelopes
+        kwargs['type'] = self.objects_to_fetch
+        if isinstance(self.collection_to_fetch, v20.Collection):
+            return v20.as_pages(get_objects, per_request=page_size, **kwargs)
+        return v21.as_pages(get_objects, per_request=page_size, **kwargs)
 
     def get_page_size(self, max_limit: int, cur_limit: int) -> int:
         """

Packs/ApiModules/Scripts/TAXII2ApiModule/TAXII2ApiModule_test.py

@@ -458,7 +458,7 @@ def test_21_empty(self):
         expected = []
         mock_client = Taxii2FeedClient(url='', collection_to_fetch='', proxies=[], verify=False, objects_to_fetch=[])
 
-        actual = mock_client.load_stix_objects_from_envelope({"indicator": STIX_ENVELOPE_NO_IOCS}, -1)
+        actual = mock_client.load_stix_objects_from_envelope(STIX_ENVELOPE_NO_IOCS, -1)
 
         assert len(actual) == 0
         assert expected == actual
@@ -481,7 +481,7 @@ def test_21_simple(self):
         mock_client = Taxii2FeedClient(url='', collection_to_fetch='', proxies=[], verify=False, tlp_color='GREEN',
                                        objects_to_fetch=[])
 
-        actual = mock_client.load_stix_objects_from_envelope({"indicator": STIX_ENVELOPE_17_IOCS_19_OBJS}, -1)
+        actual = mock_client.load_stix_objects_from_envelope(STIX_ENVELOPE_17_IOCS_19_OBJS, -1)
 
         assert len(actual) == 17
         assert expected == actual
@@ -495,7 +495,7 @@ def test_21_complex_not_skipped(self):
         - skip is False
 
         When:
-        - extract_indicators_from_envelope_and_parse is called
+        - load_stix_objects_from_envelope is called
 
         Then:
         - Extract and parse the indicators from the envelope with the complex iocs
@@ -505,7 +505,7 @@ def test_21_complex_not_skipped(self):
         mock_client = Taxii2FeedClient(url='', collection_to_fetch='', proxies=[], verify=False, tlp_color='GREEN',
                                        objects_to_fetch=[])
 
-        actual = mock_client.load_stix_objects_from_envelope({"indicator": STIX_ENVELOPE_20_IOCS_19_OBJS}, -1)
+        actual = mock_client.load_stix_objects_from_envelope(STIX_ENVELOPE_20_IOCS_19_OBJS, -1)
 
         assert len(actual) == 20
         assert actual == expected
@@ -519,7 +519,7 @@ def test_21_complex_skipped(self):
         - skip is True
 
         When:
-        - extract_indicators_from_envelope_and_parse is called
+        - load_stix_objects_from_envelope is called
 
         Then:
         - Extract and parse the indicators from the envelope with the complex iocs
@@ -529,7 +529,7 @@ def test_21_complex_skipped(self):
         mock_client = Taxii2FeedClient(url='', collection_to_fetch='', proxies=[], verify=False, skip_complex_mode=True,
                                        objects_to_fetch=[])
 
-        actual = mock_client.load_stix_objects_from_envelope({"indicator": STIX_ENVELOPE_20_IOCS_19_OBJS}, -1)
+        actual = mock_client.load_stix_objects_from_envelope(STIX_ENVELOPE_20_IOCS_19_OBJS, -1)
 
         assert len(actual) == 14
         assert actual == expected
@@ -550,7 +550,6 @@ def test_load_stix_objects_from_envelope_v21(self):
         """
         mock_client = Taxii2FeedClient(url='', collection_to_fetch='', proxies=[], verify=False, objects_to_fetch=[])
         objects_envelopes = envelopes_v21
-        mock_client.id_to_object = id_to_object
 
         result = mock_client.load_stix_objects_from_envelope(objects_envelopes, -1)
         assert mock_client.id_to_object == id_to_object
@@ -564,29 +563,15 @@ def test_load_stix_objects_from_envelope_v20(self):
         - Envelope with indicators, arranged by object type.
 
         When:
-        - parse_generator_type_envelope is called (skipping condition from load_stix_objects_from_envelope).
+        - load_stix_objects_from_envelope is called.
 
         Then: - Load and parse objects from the envelope according to their object type and ignore
         extension-definition objects.
 
         """
         mock_client = Taxii2FeedClient(url='', collection_to_fetch='', proxies=[], verify=False, objects_to_fetch=[])
-        objects_envelopes = envelopes_v20
-        mock_client.id_to_object = id_to_object
-
-        parse_stix_2_objects = {
-            "indicator": mock_client.parse_indicator,
-            "attack-pattern": mock_client.parse_attack_pattern,
-            "malware": mock_client.parse_malware,
-            "report": mock_client.parse_report,
-            "course-of-action": mock_client.parse_course_of_action,
-            "campaign": mock_client.parse_campaign,
-            "intrusion-set": mock_client.parse_intrusion_set,
-            "tool": mock_client.parse_tool,
-            "threat-actor": mock_client.parse_threat_actor,
-            "infrastructure": mock_client.parse_infrastructure
-        }
-        result = mock_client.parse_generator_type_envelope(objects_envelopes, parse_stix_2_objects)
+
+        result = mock_client.load_stix_objects_from_envelope(envelopes_v20)
         assert mock_client.id_to_object == id_to_object
         assert result == parsed_objects
 
@@ -1083,3 +1068,21 @@ def test_parse_location(self, taxii_2_client, location_object, xsoar_expected_re
          - Make sure all the fields are being parsed correctly.
         """
         assert taxii_2_client.parse_location(location_object) == xsoar_expected_response
+
+
+@pytest.mark.parametrize('limit, element_count, return_value',
+                         [(8, 8, True),
+                          (8, 9, True),
+                          (8, 0, False),
+                          (-1, 10, False)])
+def test_reached_limit(limit, element_count, return_value):
+    """
+    Given:
+        - A limit and element count.
+    When:
+        - Enforcing limit on the elements count.
+    Then:
+        - Assert that the element count is not exceeded.
+    """
+    from TAXII2ApiModule import reached_limit
+    assert reached_limit(limit, element_count) == return_value

Packs/ApiModules/Scripts/TAXII2ApiModule/test_data/id_to_object_test.json

@@ -201,23 +201,5 @@
       "modified":"2016-05-07T11:22:30.000Z",
       "name":"Poison Ivy C2",
       "infrastructure_types": ["command-and-control"]
-    },
-  "relationship--01a5a209-b94c-450b-b7f9-946497d91055": {
-      "created": "2018-08-03T21:03:51.484Z",
-      "id": "relationship--01a5a209-b94c-450b-b7f9-946497d91055",
-      "modified": "2018-08-03T21:03:51.484Z",
-      "relationship_type": "uses",
-      "source_ref": "campaign--6320584e-3ef0-4a72-aaf8-0a49fa1d477c",
-      "target_ref": "attack-pattern--4e6620ac-c30c-4f6d-918e-fa20cae7c1ce",
-      "type": "relationship"
-    },
-  "relationship--abc475d9-199c-4623-9e9a-02adf340a415": {
-      "created": "2018-08-03T20:31:03.780Z",
-      "id": "relationship--abc475d9-199c-4623-9e9a-02adf340a415",
-      "modified": "2018-08-22T12:36:32.248Z",
-      "relationship_type": "indicates",
-      "source_ref": "indicator--545928d9-bfe8-4320-bb98-751f38139892",
-      "target_ref": "campaign--6320584e-3ef0-4a72-aaf8-0a49fa1d477c",
-      "type": "relationship"
     }
 }