Merge pull request #182 from DomainTools/release-2.6.0

Release 2.6.0

Author: briluza

README.md

@@ -304,6 +304,10 @@ To add more e2e tests, put these in the `../tests/e2e` folder.
         export TEST_USER=<user-key>
         export TEST_KEY=<api-key>
     ```
+- Run unit tests.
+    ```bash
+        tox -e
+    ```
 
 ## Run the end-to-end test script
 - Before running the test, be sure that docker is running.

VERSION

@@ -1 +1 @@
-2.5.3
+2.6.0

domaintools/_version.py

@@ -20,4 +20,4 @@
 
 """
 
-current = "2.5.3"
+current = "2.6.0"

domaintools/api.py

@@ -6,7 +6,13 @@
 import re
 import ssl
 
-from domaintools.constants import Endpoint, ENDPOINT_TO_SOURCE_MAP, RTTF_PRODUCTS_LIST, OutputFormat
+from domaintools.constants import (
+    Endpoint,
+    OutputFormat,
+    ENDPOINT_TO_SOURCE_MAP,
+    RTTF_PRODUCTS_LIST,
+    RTTF_PRODUCTS_CMD_MAPPING,
+)
 from domaintools._version import current as version
 from domaintools.results import (
     GroupedIterable,
@@ -92,7 +98,9 @@ def __init__(
         self._build_api_url(api_url, api_port)
 
         if not https:
-            raise Exception("The DomainTools API endpoints no longer support http traffic. Please make sure https=True.")
+            raise Exception(
+                "The DomainTools API endpoints no longer support http traffic. Please make sure https=True."
+            )
         if proxy_url and not isinstance(proxy_url, str):
             raise Exception("Proxy URL must be a string. For example: '127.0.0.1:8888'")
 
@@ -110,8 +118,12 @@ def _build_api_url(self, api_url=None, api_port=None):
 
         self._rest_api_url = rest_api_url
 
-    def _rate_limit(self):
+    def _rate_limit(self, product):
         """Pulls in and enforces the latest rate limits for the specified user"""
+        if product in RTTF_PRODUCTS_LIST:
+            self.limits_set = False
+            return
+
         self.limits_set = True
         for product in self.account_information():
             limit_minutes = product["per_minute_limit"] or None
@@ -128,8 +140,9 @@ def _results(self, product, path, cls=Results, **kwargs):
         if product != "account-information" and self.rate_limit and not self.limits_set and not self.limits:
             always_sign_api_key_previous_value = self.always_sign_api_key
             header_authentication_previous_value = self.header_authentication
-            self._rate_limit()
-            # Reset always_sign_api_key and header_authentication to its original User-set values as these might be affected when self.account_information() was executed
+            self._rate_limit(product)
+            # Reset always_sign_api_key and header_authentication to its original
+            # User-set values as these might be affected when self.account_information() was executed
             self.always_sign_api_key = always_sign_api_key_previous_value
             self.header_authentication = header_authentication_previous_value
 
@@ -139,7 +152,13 @@ def _results(self, product, path, cls=Results, **kwargs):
         is_rttf_product = product in RTTF_PRODUCTS_LIST
         self._handle_api_key_parameters(is_rttf_product)
         self.handle_api_key(is_rttf_product, path, parameters)
-        parameters.update({key: str(value).lower() if value in (True, False) else value for key, value in kwargs.items() if value is not None})
+        parameters.update(
+            {
+                key: str(value).lower() if value in (True, False) else value
+                for key, value in kwargs.items()
+                if value is not None
+            }
+        )
 
         return cls(self, product, uri, **parameters)
 
@@ -189,8 +208,30 @@ def snakecase(string):
                 string[1:],
             )
 
-        api_calls = tuple((api_call for api_call in dir(API) if not api_call.startswith("_") and callable(getattr(API, api_call, None))))
-        return sorted([snakecase(p["id"]) for p in self.account_information()["products"] if snakecase(p["id"]) in api_calls])
+        api_calls = tuple(
+            (
+                api_call
+                for api_call in dir(API)
+                if not api_call.startswith("_") and callable(getattr(API, api_call, None))
+            )
+        )
+
+        account_information = self.account_information()
+
+        available_calls = set()
+        for product in self.account_information():
+            product_id = product["id"]
+            # for RTUF endpoints as we use different func name in our wrapper
+            if product_id in RTTF_PRODUCTS_LIST:
+                if rttf_api_command := RTTF_PRODUCTS_CMD_MAPPING.get(product_id):
+                    available_calls.add(rttf_api_command)
+
+            # for IRIS endpoints
+            snakecase_pid = snakecase(product_id)
+            if snakecase_pid in api_calls:
+                available_calls.add(snakecase_pid)
+
+        return sorted(available_calls)
 
     def brand_monitor(self, query, exclude=None, domain_status=None, days_back=None, **kwargs):
         """Pass in one or more terms as a list or separated by the pipe character ( | )"""
@@ -445,7 +486,16 @@ def iris(
         """Performs a search for the provided search terms ANDed together,
         returning the pivot engine row data for the resulting domains.
         """
-        if not domain and not ip and not email and not nameserver and not registrar and not registrant and not registrant_org and not kwargs:
+        if (
+            not domain
+            and not ip
+            and not email
+            and not nameserver
+            and not registrar
+            and not registrant
+            and not registrant_org
+            and not kwargs
+        ):
             raise ValueError("At least one search term must be specified")
 
         return self._results(
@@ -1069,7 +1119,10 @@ def nod(self, **kwargs) -> FeedsResults:
         validate_feeds_parameters(kwargs)
         endpoint = kwargs.pop("endpoint", Endpoint.FEED.value)
         source = ENDPOINT_TO_SOURCE_MAP.get(endpoint)
-        if endpoint == Endpoint.DOWNLOAD.value or kwargs.get("output_format", OutputFormat.JSONL.value) != OutputFormat.CSV.value:
+        if (
+            endpoint == Endpoint.DOWNLOAD.value
+            or kwargs.get("output_format", OutputFormat.JSONL.value) != OutputFormat.CSV.value
+        ):
             # headers param is allowed only in Feed API and CSV format
             kwargs.pop("headers", None)
 
@@ -1101,7 +1154,10 @@ def nad(self, **kwargs) -> FeedsResults:
         validate_feeds_parameters(kwargs)
         endpoint = kwargs.pop("endpoint", Endpoint.FEED.value)
         source = ENDPOINT_TO_SOURCE_MAP.get(endpoint).value
-        if endpoint == Endpoint.DOWNLOAD.value or kwargs.get("output_format", OutputFormat.JSONL.value) != OutputFormat.CSV.value:
+        if (
+            endpoint == Endpoint.DOWNLOAD.value
+            or kwargs.get("output_format", OutputFormat.JSONL.value) != OutputFormat.CSV.value
+        ):
             # headers param is allowed only in Feed API and CSV format
             kwargs.pop("headers", None)
 
@@ -1162,7 +1218,10 @@ def domaindiscovery(self, **kwargs) -> FeedsResults:
         validate_feeds_parameters(kwargs)
         endpoint = kwargs.pop("endpoint", Endpoint.FEED.value)
         source = ENDPOINT_TO_SOURCE_MAP.get(endpoint).value
-        if endpoint == Endpoint.DOWNLOAD.value or kwargs.get("output_format", OutputFormat.JSONL.value) != OutputFormat.CSV.value:
+        if (
+            endpoint == Endpoint.DOWNLOAD.value
+            or kwargs.get("output_format", OutputFormat.JSONL.value) != OutputFormat.CSV.value
+        ):
             # headers param is allowed only in Feed API and CSV format
             kwargs.pop("headers", None)
 
@@ -1194,7 +1253,10 @@ def noh(self, **kwargs) -> FeedsResults:
         validate_feeds_parameters(kwargs)
         endpoint = kwargs.pop("endpoint", Endpoint.FEED.value)
         source = ENDPOINT_TO_SOURCE_MAP.get(endpoint).value
-        if endpoint == Endpoint.DOWNLOAD.value or kwargs.get("output_format", OutputFormat.JSONL.value) != OutputFormat.CSV.value:
+        if (
+            endpoint == Endpoint.DOWNLOAD.value
+            or kwargs.get("output_format", OutputFormat.JSONL.value) != OutputFormat.CSV.value
+        ):
             # headers param is allowed only in Feed API and CSV format
             kwargs.pop("headers", None)
 
@@ -1225,12 +1287,15 @@ def realtime_domain_risk(self, **kwargs) -> FeedsResults:
         validate_feeds_parameters(kwargs)
         endpoint = kwargs.pop("endpoint", Endpoint.FEED.value)
         source = ENDPOINT_TO_SOURCE_MAP.get(endpoint).value
-        if endpoint == Endpoint.DOWNLOAD.value or kwargs.get("output_format", OutputFormat.JSONL.value) != OutputFormat.CSV.value:
+        if (
+            endpoint == Endpoint.DOWNLOAD.value
+            or kwargs.get("output_format", OutputFormat.JSONL.value) != OutputFormat.CSV.value
+        ):
             # headers param is allowed only in Feed API and CSV format
             kwargs.pop("headers", None)
 
         return self._results(
-            f"domain-risk-feed-({source})",
+            f"real-time-domain-risk-({source})",
             f"v1/{endpoint}/domainrisk/",
             response_path=(),
             cls=FeedsResults,
@@ -1256,12 +1321,15 @@ def domainhotlist(self, **kwargs) -> FeedsResults:
         validate_feeds_parameters(kwargs)
         endpoint = kwargs.pop("endpoint", Endpoint.FEED.value)
         source = ENDPOINT_TO_SOURCE_MAP.get(endpoint).value
-        if endpoint == Endpoint.DOWNLOAD.value or kwargs.get("output_format", OutputFormat.JSONL.value) != OutputFormat.CSV.value:
+        if (
+            endpoint == Endpoint.DOWNLOAD.value
+            or kwargs.get("output_format", OutputFormat.JSONL.value) != OutputFormat.CSV.value
+        ):
             # headers param is allowed only in Feed API and CSV format
             kwargs.pop("headers", None)
 
         return self._results(
-            f"domain-hotlist-feed-({source})",
+            f"real-time-domain-hotlist-({source})",
             f"v1/{endpoint}/domainhotlist/",
             response_path=(),
             cls=FeedsResults,

domaintools/base_results.py

@@ -9,7 +9,11 @@
 from datetime import datetime
 from httpx import Client
 
-from domaintools.constants import RTTF_PRODUCTS_LIST, OutputFormat, HEADER_ACCEPT_KEY_CSV_FORMAT
+from domaintools.constants import (
+    RTTF_PRODUCTS_LIST,
+    OutputFormat,
+    HEADER_ACCEPT_KEY_CSV_FORMAT,
+)
 from domaintools.exceptions import (
     BadRequestException,
     InternalServerErrorException,
@@ -53,6 +57,7 @@ def __init__(
         self._response = None
         self._items_list = None
         self._data = None
+        self._status = None
 
     def _wait_time(self):
         if not self.api.rate_limit or not self.product in self.api.limits:
@@ -92,10 +97,10 @@ def _get_session_params_and_headers(self):
             header_key_for_api_key = "X-Api-Key" if is_rttf_product else "X-API-Key"
             headers[header_key_for_api_key] = self.api.key
 
-        return {"parameters": parameters, "headers": headers}
+        session_param_and_headers = {"parameters": parameters, "headers": headers}
+        return session_param_and_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()
             headers = session_params_and_headers.get("headers")
@@ -113,7 +118,12 @@ def _make_request(self):
                 return session.patch(url=self.url, json=patch_data, headers=headers)
             else:
                 parameters = session_params_and_headers.get("parameters")
-                return session.get(url=self.url, params=parameters, headers=headers, **self.api.extra_request_params)
+                return session.get(
+                    url=self.url,
+                    params=parameters,
+                    headers=headers,
+                    **self.api.extra_request_params,
+                )
 
     def _get_results(self):
         wait_for = self._wait_time()
@@ -152,7 +162,9 @@ def data(self):
     def check_limit_exceeded(self):
         limit_exceeded, reason = False, ""
         if isinstance(self._data, dict) and (
-            "response" in self._data and "limit_exceeded" in self._data["response"] and self._data["response"]["limit_exceeded"] is True
+            "response" in self._data
+            and "limit_exceeded" in self._data["response"]
+            and self._data["response"]["limit_exceeded"] is True
         ):
             limit_exceeded, reason = True, self._data["response"]["message"]
         elif "response" in self._data and "limit_exceeded" in self._data:
@@ -163,12 +175,12 @@ def check_limit_exceeded(self):
 
     @property
     def status(self):
-        if not getattr(self, "_status", None):
+        if not getattr(self, "_status", None) and not self.product in RTTF_PRODUCTS_LIST:
             self._status = self._get_results().status_code
 
         return self._status
 
-    def setStatus(self, code, response=None):
+    def setStatus(self, code, response=None, reason_text=None):
         self._status = code
         if code == 200 or (self.product in RTTF_PRODUCTS_LIST and code == 206):
             return
@@ -181,6 +193,9 @@ def setStatus(self, code, response=None):
                 reason = response.text
                 if callable(reason):
                     reason = reason()
+        else:  # optionally pass a customize reason of error for better traceback
+            if reason_text is not None:
+                reason = reason_text
 
         if code in (400, 422):
             raise BadRequestException(code, reason)
@@ -330,4 +345,8 @@ def as_list(self):
         return "\n".join([json.dumps(item, indent=4, separators=(",", ": ")) for item in self._items()])
 
     def __str__(self):
-        return str(json.dumps(self.data(), indent=4, separators=(",", ": ")) if self.kwargs.get("format", "json") == "json" else self.data())
+        return str(
+            json.dumps(self.data(), indent=4, separators=(",", ": "))
+            if self.kwargs.get("format", "json") == "json"
+            else self.data()
+        )

domaintools/cli/api.py

@@ -58,7 +58,9 @@ def validate_after_or_before_input(value: str):
             datetime.strptime(value, "%Y-%m-%dT%H:%M:%SZ")
             return value
         except:
-            raise typer.BadParameter(f"{value} is neither an integer or a valid ISO 8601 datetime string in UTC form")
+            raise typer.BadParameter(
+                f"{value} is neither an integer or a valid ISO 8601 datetime string in UTC form"
+            )
 
     @staticmethod
     def validate_source_file_extension(value: str):
@@ -78,7 +80,9 @@ def validate_source_file_extension(value: str):
         ext = get_file_extension(value)
 
         if ext.lower() not in VALID_EXTENSIONS:
-            raise typer.BadParameter(f"{value} is not in valid extensions. Valid file extensions: {VALID_EXTENSIONS}")
+            raise typer.BadParameter(
+                f"{value} is not in valid extensions. Valid file extensions: {VALID_EXTENSIONS}"
+            )
 
         return value
 
@@ -111,7 +115,7 @@ def _get_formatted_output(cls, cmd_name: str, response, out_format: str = "json"
         if cmd_name in ("available_api_calls",):
             return "\n".join(response)
         if response.product in RTTF_PRODUCTS_LIST:
-            return "\n".join([data for data in response.response()])
+            pass  # do nothing
         return str(getattr(response, out_format) if out_format != "list" else response.as_list())
 
     @classmethod
@@ -203,7 +207,7 @@ def run(cls, name: str, params: Optional[Dict] = {}, **kwargs):
                 transient=True,
             ) as progress:
 
-                progress.add_task(
+                task_id = progress.add_task(
                     description=f"Using api credentials with a username of: [cyan]{user}[/cyan]\nExecuting [green]{name}[/green] api call...",
                     total=None,
                 )
@@ -218,27 +222,36 @@ def run(cls, name: str, params: Optional[Dict] = {}, **kwargs):
                     header_authentication=header_authentication,
                 )
                 dt_api_func = getattr(dt_api, name)
-
                 params = params | kwargs
 
                 response = dt_api_func(**params)
-                progress.add_task(
+                progress.update(
+                    task_id,
                     description=f"Preparing results with format of {response_format}...",
-                    total=None,
                 )
 
-                output = cls._get_formatted_output(cmd_name=name, response=response, out_format=response_format)
+                output = cls._get_formatted_output(
+                    cmd_name=name, response=response, out_format=response_format
+                )
 
                 if isinstance(out_file, _io.TextIOWrapper):
+                    progress.update(
+                        task_id,
+                        description=f"Printing the results with format of {response_format}...",
+                    )
                     # use rich `print` command to prettify the ouput in sys.stdout
-                    if response.product in RTTF_PRODUCTS_LIST:
-                        print(output)
+                    if name not in ("available_api_calls",) and response.product in RTTF_PRODUCTS_LIST:
+                        for feeds in response.response():
+                            print(feeds)
                     else:
                         print(response)
                 else:
+                    progress.update(
+                        task_id,
+                        description=f"Writing results to {out_file}",
+                    )
                     # if it's a file then write
                     out_file.write(output if output.endswith("\n") else output + "\n")
-                time.sleep(0.25)
         except Exception as e:
             if isinstance(e, ServiceException):
                 code = typer.style(getattr(e, "code", 400), fg=typer.colors.BRIGHT_RED)