From d649f19680a0dae223de4318e4607e0d3b96807e Mon Sep 17 00:00:00 2001
From: BK Box <bk@theboxes.org>
Date: Fri, 21 Mar 2014 15:24:28 -0500
Subject: [PATCH] Add exception `SatoriInvalidDomain`

When `domain_info` is passed an invalid domain, it would attempt to
parse the registered domain and eventually pass a blank string to
`pythonwhois.get_whois`. This causes an exception of "No root WHOIS
server found for TLD." This is a valid exception, but we should not make
it this far to get this exception.

Instead, this patch will raise an exception of `SatoriInvalidDomain` if
domain_info attempted to be used with an invalid domain.

Change-Id: Ie30b9a01b39d92d8bfeb7f3a80d333fa4bb30d49
Closes-Bug: #1295391
Closes-Bug: #1293670
---
 satori/dns.py            | 20 ++++++++++++++------
 satori/errors.py         |  5 +++++
 satori/tests/test_dns.py | 13 ++++++++++++-
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/satori/dns.py b/satori/dns.py
index 3cd1bd9..feb9831 100644
--- a/satori/dns.py
+++ b/satori/dns.py
@@ -38,11 +38,15 @@ def resolve_hostname(host):
         LOG.exception(error)
         raise errors.SatoriInvalidNetloc(error)
 
-    # Domain names are in netloc, IP addresses fall into path
+    # Domain names and IP are in netloc when parsed with a protocol
+    # they will be in path if parsed without a protocol
     hostname = parsed.netloc or parsed.path
 
-    # socket.gaierror is not trapped here
-    address = socket.gethostbyname(hostname)
+    try:
+        address = socket.gethostbyname(hostname)
+    except socket.gaierror:
+        error = "`%s` is an invalid domain." % hostname
+        raise errors.SatoriInvalidDomain(error)
     return address
 
 
@@ -53,8 +57,12 @@ def get_registered_domain(hostname):
 
 def domain_info(domain):
     """Get as much information as possible for a given domain name."""
-    domain = get_registered_domain(domain)
-    result = pythonwhois.get_whois(domain)
+    registered_domain = get_registered_domain(domain)
+    if utils.is_valid_ip_address(domain) or registered_domain == '':
+        error = "`%s` is an invalid domain." % domain
+        raise errors.SatoriInvalidDomain(error)
+
+    result = pythonwhois.get_whois(registered_domain)
     registrar = []
     if 'registrar' in result and len(result['registrar']) > 0:
         registrar = result['registrar'][0]
@@ -72,7 +80,7 @@ def domain_info(domain):
                 days_until_expires = (utils.parse_time_string(expires) -
                                       datetime.datetime.now()).days
     return {
-        'name': domain,
+        'name': registered_domain,
         'whois': result['raw'],
         'registrar': registrar,
         'nameservers': nameservers,
diff --git a/satori/errors.py b/satori/errors.py
index 314ca3e..9b53b49 100644
--- a/satori/errors.py
+++ b/satori/errors.py
@@ -26,6 +26,11 @@ class SatoriInvalidNetloc(SatoriException):
     """Netloc that cannot be parsed by `urlparse`."""
 
 
+class SatoriInvalidDomain(SatoriException):
+
+    """Invalid Domain provided."""
+
+
 class SatoriShellException(SatoriException):
 
     """Invalid shell parameters."""
diff --git a/satori/tests/test_dns.py b/satori/tests/test_dns.py
index 8af2104..5bf4962 100644
--- a/satori/tests/test_dns.py
+++ b/satori/tests/test_dns.py
@@ -12,7 +12,6 @@
 #
 """Satori DNS Discovery."""
 
-import datetime
 import socket
 import unittest
 
@@ -265,5 +264,17 @@ class TestDNS(utils.TestCase):
         data = dns.domain_info(self.domain)
         self.assertIsNone(data['expiration_date'])
 
+    def test_domain_info_raises_invalid_domain_error(self):
+        ip_whois = ["""
+        Home net HOME-NET-192-168 (NET-192-0-0-0-1)
+        Home Inc. HOME-NET-192-168-0 (NET-192-168-0-0-1)
+        """]
+        self.mock_get_whois_raw.return_value = ip_whois
+        self.assertRaises(
+            errors.SatoriInvalidDomain,
+            dns.domain_info,
+            "192.168.0.1"
+        )
+
 if __name__ == "__main__":
     unittest.main()