Add support for GeoLite2 ASN database

Author: oschwald

HISTORY.rst

@@ -3,6 +3,11 @@
 History
 -------
 
+2.5.0
+++++++++++++++++++
+
+* Added support for GeoLite2 ASN database.
+
 2.4.2 (2016-12-08)
 ++++++++++++++++++
 

README.rst

@@ -185,6 +185,22 @@ Anonymous IP Database
     '128.101.101.101'
     >>> reader.close()
 
+ASN Database
+^^^^^^^^^^^^
+
+.. code-block:: pycon
+
+    >>> import geoip2.database
+    >>>
+    >>> # This creates a Reader object. You should use the same object
+    >>> # across multiple requests as creation of it is expensive.
+    >>> with geoip2.database.Reader('/path/to/GeoLite2-ASN.mmdb') as reader:
+    >>>     response = reader.asn('1.128.0.0')
+    >>>     response.autonomous_system_number
+    1221
+    >>>     response.autonomous_system_organization
+    'Telstra Pty Ltd'
+
 Connection-Type Database
 ^^^^^^^^^^^^^^^^^^^^^^^^
 

geoip2/database.py

@@ -120,6 +120,17 @@ def anonymous_ip(self, ip_address):
         return self._flat_model_for(geoip2.models.AnonymousIP,
                                     'GeoIP2-Anonymous-IP', ip_address)
 
+    def asn(self, ip_address):
+        """Get the ASN object for the IP address.
+
+        :param ip_address: IPv4 or IPv6 address as a string.
+
+        :returns: :py:class:`geoip2.models.ASN` object
+
+        """
+        return self._flat_model_for(geoip2.models.ASN, 'GeoLite2-ASN',
+                                    ip_address)
+
     def connection_type(self, ip_address):
         """Get the ConnectionType object for the IP address.
 

geoip2/models.py

@@ -368,6 +368,40 @@ def __init__(self, raw):
         self.raw = raw
 
 
+class ASN(SimpleModel):
+    """Model class for the GeoLite2 ASN.
+
+    This class provides the following attribute:
+
+    .. attribute:: autonomous_system_number
+
+      The autonomous system number associated with the IP address.
+
+      :type: int
+
+    .. attribute:: autonomous_system_organization
+
+      The organization associated with the registered autonomous system number
+      for the IP address.
+
+      :type: unicode
+
+    .. attribute:: ip_address
+
+      The IP address used in the lookup.
+
+      :type: unicode
+    """
+
+    # pylint:disable=too-many-arguments
+    def __init__(self, raw):
+        self.autonomous_system_number = raw.get('autonomous_system_number')
+        self.autonomous_system_organization = raw.get(
+            'autonomous_system_organization')
+        self.ip_address = raw.get('ip_address')
+        self.raw = raw
+
+
 class ConnectionType(SimpleModel):
     """Model class for the GeoIP2 Connection-Type.
 
@@ -424,7 +458,7 @@ def __init__(self, raw):
         self.raw = raw
 
 
-class ISP(SimpleModel):
+class ISP(ASN):
     """Model class for the GeoIP2 ISP.
 
     This class provides the following attribute:
@@ -463,10 +497,6 @@ class ISP(SimpleModel):
 
     # pylint:disable=too-many-arguments
     def __init__(self, raw):
-        self.autonomous_system_number = raw.get('autonomous_system_number')
-        self.autonomous_system_organization = raw.get(
-            'autonomous_system_organization')
+        super(ISP, self).__init__(raw)
         self.isp = raw.get('isp')
         self.organization = raw.get('organization')
-        self.ip_address = raw.get('ip_address')
-        self.raw = raw

tests/data

@@ -76,6 +76,26 @@ def test_anonymous_ip(self):
         self.assertEqual(record.ip_address, ip_address)
         reader.close()
 
+    def test_asn(self):
+        reader = geoip2.database.Reader(
+            'tests/data/test-data/GeoLite2-ASN-Test.mmdb')
+
+        ip_address = '1.128.0.0'
+        record = reader.asn(ip_address)
+        self.assertEqual(record.autonomous_system_number, 1221)
+        self.assertEqual(record.autonomous_system_organization,
+                         'Telstra Pty Ltd')
+        self.assertEqual(record.ip_address, ip_address)
+
+        self.assertRegex(
+            str(record),
+            r'geoip2.models.ASN\(.*1\.128\.0\.0.*\)',
+            'str representation is correct')
+
+        self.assertEqual(record, eval(repr(record)), "ASN repr can be eval'd")
+
+        reader.close()
+
     def test_city(self):
         reader = geoip2.database.Reader(
             'tests/data/test-data/GeoIP2-City-Test.mmdb')