IDEV-2020: Use the FeedsResults class.

Author: jbabac

domaintools/api.py

@@ -1,6 +1,8 @@
 from datetime import datetime, timedelta, timezone
 from hashlib import sha1, sha256, md5
 from hmac import new as hmac
+from types import GeneratorType
+
 import re
 
 from domaintools.constants import Endpoint, ENDPOINT_TO_SOURCE_MAP, FEEDS_PRODUCTS_LIST, OutputFormat
@@ -11,6 +13,7 @@
     ParsedDomainRdap,
     Reputation,
     Results,
+    FeedsResults,
 )
 from domaintools.filters import (
     filter_by_riskscore,
@@ -1065,7 +1068,7 @@ def iris_detect_ignored_domains(
             **kwargs,
         )
 
-    def nod(self, **kwargs):
+    def nod(self, **kwargs) -> GeneratorType:
         """Returns back list of the newly observed domains feed"""
         validate_feeds_parameters(kwargs)
         endpoint = kwargs.pop("endpoint", Endpoint.FEED.value)
@@ -1078,10 +1081,11 @@ def nod(self, **kwargs):
             f"newly-observed-domains-feed-({source.value})",
             f"v1/{endpoint}/nod/",
             response_path=(),
+            cls=FeedsResults,
             **kwargs,
         )
 
-    def nad(self, **kwargs):
+    def nad(self, **kwargs) -> GeneratorType:
         """Returns back list of the newly active domains feed"""
         validate_feeds_parameters(kwargs)
         endpoint = kwargs.pop("endpoint", Endpoint.FEED.value)
@@ -1094,10 +1098,11 @@ def nad(self, **kwargs):
             f"newly-active-domains-feed-({source})",
             f"v1/{endpoint}/nad/",
             response_path=(),
+            cls=FeedsResults,
             **kwargs,
         )
 
-    def domainrdap(self, **kwargs):
+    def domainrdap(self, **kwargs) -> GeneratorType:
         """Returns changes to global domain registration information, populated by the Registration Data Access Protocol (RDAP)"""
         validate_feeds_parameters(kwargs)
         endpoint = kwargs.pop("endpoint", Endpoint.FEED.value)
@@ -1107,10 +1112,11 @@ def domainrdap(self, **kwargs):
             f"domain-registration-data-access-protocol-feed-({source})",
             f"v1/{endpoint}/domainrdap/",
             response_path=(),
+            cls=FeedsResults,
             **kwargs,
         )
 
-    def domaindiscovery(self, **kwargs):
+    def domaindiscovery(self, **kwargs) -> GeneratorType:
         """Returns new domains as they are either discovered in domain registration information, observed by our global sensor network, or reported by trusted third parties"""
         validate_feeds_parameters(kwargs)
         endpoint = kwargs.pop("endpoint", Endpoint.FEED.value)
@@ -1123,5 +1129,6 @@ def domaindiscovery(self, **kwargs):
             f"real-time-domain-discovery-feed-({source})",
             f"v1/{endpoint}/domaindiscovery/",
             response_path=(),
+            cls=FeedsResults,
             **kwargs,
         )

domaintools/base_results.py

@@ -77,6 +77,23 @@ def _wait_time(self):
 
         return wait_for
 
+    def _get_session_params(self):
+        parameters = deepcopy(self.kwargs)
+        parameters.pop("output_format", None)
+        parameters.pop(
+            "format", None
+        )  # For some unknownn reasons, even if "format" is not included in the cli params for feeds endpoint, it is being populated thus we need to remove it. Happens only if using CLI.
+        headers = {}
+        if self.kwargs.get("output_format", OutputFormat.JSONL.value) == OutputFormat.CSV.value:
+            parameters["headers"] = int(bool(self.kwargs.get("headers", False)))
+            headers["accept"] = HEADER_ACCEPT_KEY_CSV_FORMAT
+
+        header_api_key = parameters.pop("X-Api-Key", None)
+        if header_api_key:
+            headers["X-Api-Key"] = header_api_key
+
+        return {"parameters": parameters, "headers": headers}
+
     def _make_request(self):
 
         with Client(verify=self.api.verify_ssl, proxy=self.api.proxy_url, timeout=None) as session:
@@ -93,20 +110,9 @@ def _make_request(self):
                 patch_data.update(self.api.extra_request_params)
                 return session.patch(url=self.url, json=patch_data)
             elif self.product in FEEDS_PRODUCTS_LIST:
-                parameters = deepcopy(self.kwargs)
-                parameters.pop("output_format", None)
-                parameters.pop(
-                    "format", None
-                )  # For some unknownn reasons, even if "format" is not included in the cli params for feeds endpoint, it is being populated thus we need to remove it. Happens only if using CLI.
-                headers = {}
-                if self.kwargs.get("output_format", OutputFormat.JSONL.value) == OutputFormat.CSV.value:
-                    parameters["headers"] = int(bool(self.kwargs.get("headers", False)))
-                    headers["accept"] = HEADER_ACCEPT_KEY_CSV_FORMAT
-
-                header_api_key = parameters.pop("X-Api-Key", None)
-                if header_api_key:
-                    headers["X-Api-Key"] = header_api_key
-
+                session_params = self._get_session_params()
+                parameters = session_params.get("parameters")
+                headers = session_params.get("headers")
                 return session.get(url=self.url, params=parameters, headers=headers, **self.api.extra_request_params)
             else:
                 return session.get(url=self.url, params=self.kwargs, **self.api.extra_request_params)
@@ -158,7 +164,6 @@ def check_limit_exceeded(self):
         if self.kwargs.get("format", "json") == "json" and self.product not in FEEDS_PRODUCTS_LIST:
             if "response" in self._data and "limit_exceeded" in self._data["response"] and self._data["response"]["limit_exceeded"] is True:
                 return True, self._data["response"]["message"]
-        # TODO: handle html, xml response errors better.
         elif "response" in self._data and "limit_exceeded" in self._data:
             return True, "limit exceeded"
         return False, ""
@@ -172,7 +177,7 @@ def status(self):
 
     def setStatus(self, code, response=None):
         self._status = code
-        if code == 200:
+        if code == 200 or (self.product in FEEDS_PRODUCTS_LIST and code == 206):
             return
 
         reason = None