Improvements to !geolocate (earwig/earwigbot#41).

commands/geolocate.py

@@ -20,8 +20,10 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-import json
-import urllib2
+from json import loads
+from socket import (AF_INET, AF_INET6, error as socket_error, gethostbyname,
+                    inet_pton)
+from urllib2 import urlopen
 
 from earwigbot.commands import Command
 
@@ -40,10 +42,6 @@ class Geolocate(Command):
             self.logger.warn(log.format(self.name))
 
     def process(self, data):
-        if not data.args:
-            self.reply(data, "Please specify an IP to lookup.")
-            return
-
         if not self.key:
             msg = 'I need an API key for http://ipinfodb.com/ stored as \x0303config.commands["{0}"]["apiKey"]\x0F.'
             log = 'Need an API key for http://ipinfodb.com/ stored as config.commands["{0}"]["apiKey"]'
@@ -51,10 +49,24 @@ class Geolocate(Command):
             self.logger.error(log.format(self.name))
             return
 
+        if data.args:
+            address = data.args[0]
+        else:
+            try:
+                address = gethostbyname(data.host)
+            except socket_error:
+                msg = "your hostname, \x0302{0}\x0F, is not an IP address!"
+                self.reply(data, msg.format(data.host))
+                return
         address = data.args[0]
+        if not self.is_ip(address):
+            msg = "\x0302{0}\x0F is not an IP address!"
+            self.reply(data, msg.format(address))
+            return
+
         url = "http://api.ipinfodb.com/v3/ip-city/?key={0}&ip={1}&format=json"
-        query = urllib2.urlopen(url.format(self.key, address)).read()
-        res = json.loads(query)
+        query = urlopen(url.format(self.key, address)).read()
+        res = loads(query)
 
         country = res["countryName"].title()
         region = res["regionName"].title()
@@ -72,3 +84,17 @@ class Geolocate(Command):
         msg = "{0}, {1}, {2} ({3}, {4}), UTC {5}"
         geo = msg.format(country, region, city, latitude, longitude, utcoffset)
         self.reply(data, geo)
+
+    def is_ip(self, address):
+        """Return ``True`` if the input is an IP address, else ``False``.
+
+        This tests for IPv4 and IPv6 using :py:func:`socket.inet_pton`.
+        """
+        try:
+            inet_pton(AF_INET, address)
+        except socket_error:
+            try:
+                inet_pton(AF_INET6, address)
+            except socket_error:
+                return False
+        return True