Cleanup code and add improvements in the implementation.

Author: jbabac

domaintools/api.py

@@ -147,15 +147,12 @@ def _handle_api_key_parameters(self, is_rttf_product):
 
     def handle_api_key(self, is_rttf_product, path, parameters):
         if self.https and not self.always_sign_api_key:
-            if self.header_authentication:
-                parameters["X-Api-Key"] = self.key
-            else:
-                parameters["api_key"] = self.key
+            parameters["api_key"] = self.key
         else:
             if is_rttf_product:
                 # As per requirement in IDEV-2272, raise this error when the user explicitly sets signing of API key for RTTF endpoints
                 raise ValueError("Real Time Threat Feeds do not support signed API keys.")
-            elif self.key_sign_hash and self.key_sign_hash in AVAILABLE_KEY_SIGN_HASHES:
+            if self.key_sign_hash and self.key_sign_hash in AVAILABLE_KEY_SIGN_HASHES:
                 signing_hash = eval(self.key_sign_hash)
             else:
                 raise ValueError(

domaintools/base_results.py

@@ -75,45 +75,45 @@ 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.
+    def _get_session_params_and_headers_and_headers(self):
         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
+        parameters = deepcopy(self.kwargs)
+        is_rttf_product = self.product in RTTF_PRODUCTS_LIST
+        if is_rttf_product:
+            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.
+            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
+
+        if self.api.header_authentication:
+            header_key_for_api_key = "X-Api-Key" if is_rttf_product else "X-API-Key"
+            headers[header_key_for_api_key] = parameters.pop("api_key", None)
 
         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:
+            session_params_and_headers = self._get_session_params_and_headers_and_headers()
+            headers = session_params_and_headers.get("headers")
             if self.product in [
                 "iris-investigate",
                 "iris-enrich",
                 "iris-detect-escalate-domains",
             ]:
                 post_data = self.kwargs.copy()
                 post_data.update(self.api.extra_request_params)
-                return session.post(url=self.url, data=post_data)
+                return session.post(url=self.url, data=post_data, headers=headers)
             elif self.product in ["iris-detect-manage-watchlist-domains"]:
                 patch_data = self.kwargs.copy()
                 patch_data.update(self.api.extra_request_params)
-                return session.patch(url=self.url, json=patch_data)
-            elif self.product in RTTF_PRODUCTS_LIST:
-                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)
+                return session.patch(url=self.url, json=patch_data, headers=headers)
             else:
-                return session.get(url=self.url, params=self.kwargs, **self.api.extra_request_params)
+                parameters = session_params_and_headers.get("parameters")
+                return session.get(url=self.url, params=parameters, headers=headers, **self.api.extra_request_params)
 
     def _get_results(self):
         wait_for = self._wait_time()

domaintools_async/__init__.py

@@ -43,21 +43,19 @@ def __await__(self):
         return self.__awaitable__().__await__()
 
     async def _make_async_request(self, session):
+        session_params_and_headers = self._get_session_params_and_headers_and_headers()
+        headers = session_params_and_headers.get("headers")
         if self.product in ["iris-investigate", "iris-enrich", "iris-detect-escalate-domains"]:
             post_data = self.kwargs.copy()
             post_data.update(self.api.extra_request_params)
-            results = await session.post(url=self.url, data=post_data)
+            results = await session.post(url=self.url, data=post_data, headers=headers)
         elif self.product in ["iris-detect-manage-watchlist-domains"]:
             patch_data = self.kwargs.copy()
-            patch_data.update(self.api.extra_request_params)
+            patch_data.update(self.api.extra_request_params, headers=headers)
             results = await session.patch(url=self.url, json=patch_data)
-        elif self.product in RTTF_PRODUCTS_LIST:
-            session_params = self._get_session_params()
-            parameters = session_params.get("parameters")
-            headers = session_params.get("headers")
-            results = await session.get(url=self.url, params=parameters, headers=headers, **self.api.extra_request_params)
         else:
-            results = await session.get(url=self.url, params=self.kwargs, **self.api.extra_request_params)
+            parameters = session_params_and_headers.get("parameters")
+            results = await session.get(url=self.url, params=parameters, headers=headers, **self.api.extra_request_params)
         if results:
             self.setStatus(results.status_code, results)
             if self.kwargs.get("format", "json") == "json":