max_depth implementation

max_depth option for whois and ai_whois queries.

asyncwhois/__init__.py

@@ -52,6 +52,7 @@ def whois(
     proxy_url: str = None,
     timeout: int = 10,
     tldextract_obj: TLDExtract = None,
+    max_depth: int = None,
 ) -> tuple[str, dict]:
     """
     Performs a WHOIS query for the given `search_term`. If `search_term` is or can be cast to an
@@ -89,6 +90,7 @@ def whois(
                 proxy_url=proxy_url,
                 timeout=timeout,
                 tldextract_obj=tldextract_obj,
+                max_depth=max_depth,
             ).whois(search_term)
     else:
         return "", {}
@@ -145,6 +147,7 @@ async def aio_whois(
     proxy_url: str = None,
     timeout: int = 10,
     tldextract_obj: TLDExtract = None,
+    max_depth: int = None,
 ) -> tuple[str, dict]:
     """
     Performs a WHOIS query for the given `search_term`. If `search_term` is or can be cast to an
@@ -182,6 +185,7 @@ async def aio_whois(
                 proxy_url=proxy_url,
                 timeout=timeout,
                 tldextract_obj=tldextract_obj,
+                max_depth=max_depth,
             ).aio_whois(search_term)
     else:
         return "", {}

asyncwhois/client.py

@@ -54,14 +54,15 @@ def __init__(
         whodap_client: whodap.DNSClient = None,
         timeout: int = 10,
         tldextract_obj: TLDExtract = None,
+        max_depth: int = None,
     ):
         super().__init__(whodap_client)
         self.authoritative_only = authoritative_only
         self.ignore_not_found = ignore_not_found
         self.proxy_url = proxy_url
         self.timeout = timeout
         self.tldextract_obj = tldextract_obj
-        self.query_obj = DomainQuery(proxy_url=proxy_url, timeout=timeout)
+        self.query_obj = DomainQuery(proxy_url=proxy_url, timeout=timeout, max_depth=max_depth)
         self.parse_obj = DomainParser(ignore_not_found=ignore_not_found)
 
     def _get_domain_components(self, domain: str) -> tuple[str, str, str]:

asyncwhois/query.py

@@ -19,10 +19,10 @@ class Query:
     refer_regex = r"refer: *(.+)"
     whois_server_regex = r".+ whois server: *(.+)"
 
-    def __init__(self, proxy_url: str = None, timeout: int = 10):
+    def __init__(self, proxy_url: str = None, timeout: int = 10, max_depth: int = None):
         self.proxy_url = proxy_url
         self.timeout = timeout
-
+        self.max_depth = max_depth
     @staticmethod
     def _find_match(regex: str, blob: str) -> str:
         match = ""
@@ -130,7 +130,7 @@ async def aio_run(self, search_term: str, server: str = None) -> list[str]:
         return await self._aio_do_query(server, data, server_regex, [])
 
     def _do_query(
-        self, server: str, data: str, regex: str, chain: list[str]
+        self, server: str, data: str, regex: str, chain: list[str], depth: int = 0
     ) -> list[str]:
         """
         Recursively submits WHOIS queries until it reaches the Authoritative Server.
@@ -141,6 +141,9 @@ def _do_query(
             query_output = self._send_and_recv(conn, data)
             # save query chain
             chain.append(query_output)
+            # if max depth is reached, return the chain
+            if self.max_depth and depth >= self.max_depth:
+                return chain            
             # parse response for the referred WHOIS server name
             whois_server = self._find_match(regex, query_output)
             whois_server = whois_server.lower()
@@ -152,13 +155,13 @@ def _do_query(
             ):
                 # recursive call to find more authoritative server
                 chain = self._do_query(
-                    whois_server, data, self.whois_server_regex, chain
+                    whois_server, data, self.whois_server_regex, chain, depth + 1
                 )
         # return the WHOIS query chain
         return chain
 
     async def _aio_do_query(
-        self, server: str, data: str, regex: str, chain: list[str]
+        self, server: str, data: str, regex: str, chain: list[str], depth: int = 0
     ) -> list[str]:
         # connect to whois://<server>:43
         async with self._aio_create_connection(
@@ -171,6 +174,9 @@ async def _aio_do_query(
                 self._aio_send_and_recv(reader, writer, data), self.timeout
             )
             chain.append(query_output)
+            # if max depth is reached, return the chain
+            if self.max_depth is not None and depth >= self.max_depth:
+                return chain
             # parse response for the referred WHOIS server name
             whois_server = self._find_match(regex, query_output)
             whois_server = whois_server.lower()
@@ -183,7 +189,7 @@ async def _aio_do_query(
             ):
                 # recursive call to find the authoritative server
                 chain = await self._aio_do_query(
-                    whois_server, data, self.whois_server_regex, chain
+                    whois_server, data, self.whois_server_regex, chain, depth + 1
                 )
         # return the WHOIS query chain
         return chain
@@ -195,8 +201,9 @@ def __init__(
         server: str = None,
         proxy_url: str = None,
         timeout: int = 10,
+        max_depth: int = None,
     ):
-        super().__init__(proxy_url, timeout)
+        super().__init__(proxy_url, timeout, max_depth)
         self.server = server
 
     @staticmethod