Merge pull request #279 from iakov-kaiumov/socket-error

Raise an exception on socket error

Author: richardpenman

whois/__init__.py

@@ -33,6 +33,7 @@ def whois(
     executable_opts: Optional[list[str]] = None,
     inc_raw: bool = False,
     quiet: bool = False,
+    ignore_socket_errors: bool = True,
     convert_punycode: bool = True,
 ) -> dict[str, Any]:
     """
@@ -42,6 +43,7 @@ def whois(
     flags: flags to pass to the whois client (default 0)
     inc_raw: whether to include the raw text from whois in the result (default False)
     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)
     """
     # clean domain to expose netloc
@@ -74,7 +76,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)
+        text = nic_client.whois_lookup(None, domain, flags, quiet=quiet, ignore_socket_errors=ignore_socket_errors)
     entry = WhoisEntry.load(domain, text)
     if inc_raw:
         entry["raw"] = text

whois/whois.py

@@ -201,14 +201,17 @@ def whois(
         many_results: bool = False,
         quiet: bool = False,
         timeout: int = 10,
+        ignore_socket_errors: bool = True
     ) -> str:
         """Perform initial lookup with TLD whois server
         then, if the quick flag is false, search that result
         for the region-specific whois server and do a lookup
         there for contact details.  If `quiet` is `True`, will
         not send a message to logger when a socket error
         is encountered. Uses `timeout` as a number of seconds
-        to set as a timeout on the socket
+        to set as a timeout on the socket. If `ignore_socket_errors`
+        is `False`, will raise an exception instead of returning
+        a string containing the error.
         """
         response = b""
         s = self.get_socket()
@@ -233,24 +236,27 @@ def whois(
                 response += d
                 if not d:
                     break
-            s.close()
 
             nhost = None
             response_str = response.decode("utf-8", "replace")
             if 'with "=xxx"' in response_str:
-                return self.whois(query, hostname, flags, True)
+                return self.whois(query, hostname, flags, True, quiet=quiet, ignore_socket_errors=ignore_socket_errors)
             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=True)
+                response_str += self.whois(query, nhost, 0, quiet=quiet, ignore_socket_errors=ignore_socket_errors)
         except socket.error as e:
-            # 'response' is assigned a value (also a str) even on socket timeout
             if not quiet:
                 logger.error(
                     "Error trying to connect to socket: closing socket - {}".format(e)
                 )
+            if ignore_socket_errors:
+                # 'response' is assigned a value (also a str) even on socket timeout
+                response_str = "Socket not responding: {}".format(e)
+            else:
+                raise e
+        finally:
             s.close()
-            response_str = "Socket not responding: {}".format(e)
         return response_str
 
     def choose_server(self, domain: str) -> Optional[str]:
@@ -407,13 +413,15 @@ 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
+        self, options: Optional[dict], query_arg: str, flags: int, quiet: bool = False, ignore_socket_errors: bool = True
     ) -> str:
         """Main entry point: Perform initial lookup on TLD whois server,
         or other server to get region-specific whois server, then if quick
         flag is false, perform a second lookup on the region-specific
         server for contact records.  If `quiet` is `True`, no message
-        will be printed to STDOUT when a socket error is encountered."""
+        will be printed to STDOUT when a socket error is encountered.
+        If `ignore_socket_errors` is `False`, will raise an exception
+        instead of returning a string containing the error."""
         nichost = None
         # whoud happen when this function is called by other than main
         if options is None:
@@ -433,15 +441,16 @@ def whois_lookup(
                 options["country"] + NICClient.QNICHOST_TAIL,
                 flags,
                 quiet=quiet,
+                ignore_socket_errors=ignore_socket_errors
             )
         elif self.use_qnichost:
             nichost = self.choose_server(query_arg)
             if nichost is not None:
-                result = self.whois(query_arg, nichost, flags, quiet=quiet)
+                result = self.whois(query_arg, nichost, flags, quiet=quiet, ignore_socket_errors=ignore_socket_errors)
             else:
                 result = ""
         else:
-            result = self.whois(query_arg, options["whoishost"], flags, quiet=quiet)
+            result = self.whois(query_arg, options["whoishost"], flags, quiet=quiet, ignore_socket_errors=ignore_socket_errors)
         return result