Merge pull request #295 from goodlandsecurity/master

Allow setting timeout for WHOIS request, regex parse additional WHOIS fields

Author: richardpenman

whois/__init__.py

@@ -30,6 +30,7 @@ def whois(
     quiet: bool = False,
     ignore_socket_errors: bool = True,
     convert_punycode: bool = True,
+    timeout: int = 10,
 ) -> dict[str, Any]:
     """
     url: the URL to search whois
@@ -40,6 +41,7 @@ def whois(
     quiet: whether to avoid printing output (default False)
     ignore_socket_errors: whether to ignore socket errors (default True)
     convert_punycode: whether to convert the given URL punycode (default True)
+    timeout: timeout for WHOIS request (default 10 seconds)
     """
     # clean domain to expose netloc
     ip_match = IPV4_OR_V6.match(url)
@@ -71,7 +73,7 @@ def whois(
         nic_client = NICClient()
         if convert_punycode:
             domain = domain.encode("idna").decode("utf-8")
-        text = nic_client.whois_lookup(None, domain, flags, quiet=quiet, ignore_socket_errors=ignore_socket_errors)
+        text = nic_client.whois_lookup(None, domain, flags, quiet=quiet, ignore_socket_errors=ignore_socket_errors, timeout=timeout)
     entry = WhoisEntry.load(domain, text)
     if inc_raw:
         entry["raw"] = text

whois/parser.py

@@ -121,6 +121,10 @@ class WhoisEntry(dict):
         "state": r"Registrant State/Province: *(.+)",
         "registrant_postal_code": r"Registrant Postal Code: *(.+)",
         "country": r"Registrant Country: *(.+)",
+        "tech_name": r"Tech Name: *(.+)",
+        "tech_org": r"Tech Organization: *(.+)",
+        "admin_name": r"Admin Name: *(.+)",
+        "admin_org": r"Admin Organization: *(.+)"
     }
 
     # allows for data string manipulation before casting to date

whois/whois.py

@@ -241,11 +241,11 @@ def whois(
             nhost = None
             response_str = response.decode("utf-8", "replace")
             if 'with "=xxx"' in response_str:
-                return self.whois(query, hostname, flags, True, quiet=quiet, ignore_socket_errors=ignore_socket_errors)
+                return self.whois(query, hostname, flags, True, quiet=quiet, ignore_socket_errors=ignore_socket_errors, timeout=timeout)
             if flags & NICClient.WHOIS_RECURSE and nhost is None:
                 nhost = self.findwhois_server(response_str, hostname, query)
             if nhost is not None and nhost != "":
-                response_str += self.whois(query, nhost, 0, quiet=quiet, ignore_socket_errors=ignore_socket_errors)
+                response_str += self.whois(query, nhost, 0, quiet=quiet, ignore_socket_errors=ignore_socket_errors, timeout=timeout)
         except socket.error as e:
             if not quiet:
                 logger.error(
@@ -416,7 +416,7 @@ def choose_server(self, domain: str) -> Optional[str]:
             # return server
 
     def whois_lookup(
-        self, options: Optional[dict], query_arg: str, flags: int, quiet: bool = False, ignore_socket_errors: bool = True
+        self, options: Optional[dict], query_arg: str, flags: int, quiet: bool = False, ignore_socket_errors: bool = True, timeout: int = 10
     ) -> str:
         """Main entry point: Perform initial lookup on TLD whois server,
         or other server to get region-specific whois server, then if quick
@@ -444,16 +444,17 @@ def whois_lookup(
                 options["country"] + NICClient.QNICHOST_TAIL,
                 flags,
                 quiet=quiet,
-                ignore_socket_errors=ignore_socket_errors
+                ignore_socket_errors=ignore_socket_errors,
+                timeout=timeout
             )
         elif self.use_qnichost:
             nichost = self.choose_server(query_arg)
             if nichost is not None:
-                result = self.whois(query_arg, nichost, flags, quiet=quiet, ignore_socket_errors=ignore_socket_errors)
+                result = self.whois(query_arg, nichost, flags, quiet=quiet, ignore_socket_errors=ignore_socket_errors, timeout=timeout)
             else:
                 result = ""
         else:
-            result = self.whois(query_arg, options["whoishost"], flags, quiet=quiet, ignore_socket_errors=ignore_socket_errors)
+            result = self.whois(query_arg, options["whoishost"], flags, quiet=quiet, ignore_socket_errors=ignore_socket_errors, timeout=timeout)
         return result
 
 
@@ -599,6 +600,14 @@ def parse_command_line(argv: list[str]) -> tuple[optparse.Values, list[str]]:
         dest="whoishost",
         help="Lookup using host " + NICClient.PANDIHOST,
     )
+    parser.add_option(
+        "-t",
+        "--timeout",
+        action="store",
+        type="int",
+        dest="timeout",
+        help="Set timeout for WHOIS request",
+    )
     parser.add_option("-?", "--help", action="help")
 
     return parser.parse_args(argv)