Merge pull request #50 from maxmind/greg/ipaddress-get-support

Add support for ipaddress objects in "get" method

Author: horgh

HISTORY.rst

@@ -10,6 +10,10 @@ History
 * The extension source directory was moved to prevent an ``ImportWarning``
   when importing the module on Python 2 with ``-Wdefault`` set. Reported by
   David Szotten and Craig de Stigter. GitHub #31.
+* The ``get`` method now accepts ``ipaddress.IPv4Address`` and
+  ``ipaddress.IPv6Address`` objects in addition to strings.  This works with
+  both the pure Python implementation as well as the extension. Based on a
+  pull request #48 by Eric Pruitt. GitHub #50.
 
 1.4.1 (2018-06-22)
 ++++++++++++++++++

extension/maxminddb.c

@@ -1,5 +1,8 @@
 #include <Python.h>
 #include <maxminddb.h>
+#include <arpa/inet.h>
+#include <netinet/in.h>
+#include <sys/socket.h>
 #include "structmember.h"
 
 #define __STDC_FORMAT_MACROS
@@ -28,10 +31,12 @@ typedef struct {
     PyObject *record_size;
 } Metadata_obj;
 
+static bool format_sockaddr(struct sockaddr *addr, char *dst);
 static PyObject *from_entry_data_list(MMDB_entry_data_list_s **entry_data_list);
 static PyObject *from_map(MMDB_entry_data_list_s **entry_data_list);
 static PyObject *from_array(MMDB_entry_data_list_s **entry_data_list);
 static PyObject *from_uint128(const MMDB_entry_data_list_s *entry_data_list);
+static int ip_converter(PyObject *obj, struct sockaddr_storage *ip_address);
 
 #if PY_MAJOR_VERSION >= 3
     #define MOD_INIT(name) PyMODINIT_FUNC PyInit_ ## name(void)
@@ -106,33 +111,29 @@ static int Reader_init(PyObject *self, PyObject *args, PyObject *kwds)
 
 static PyObject *Reader_get(PyObject *self, PyObject *args)
 {
-    char *ip_address = NULL;
+    MMDB_s *mmdb = ((Reader_obj *)self)->mmdb;
 
-    Reader_obj *mmdb_obj = (Reader_obj *)self;
-    if (!PyArg_ParseTuple(args, "s", &ip_address)) {
+    if (NULL == mmdb) {
+        PyErr_SetString(PyExc_ValueError,
+                        "Attempt to read from a closed MaxMind DB.");
         return NULL;
     }
 
-    MMDB_s *mmdb = mmdb_obj->mmdb;
+    struct sockaddr_storage ip_address_ss = { 0 };
+    struct sockaddr *ip_address = (struct sockaddr *)&ip_address_ss;
+    if (!PyArg_ParseTuple(args, "O&", ip_converter, &ip_address_ss)) {
+        return NULL;
+    }
 
-    if (NULL == mmdb) {
+    if (!ip_address->sa_family) {
         PyErr_SetString(PyExc_ValueError,
-                        "Attempt to read from a closed MaxMind DB.");
+                        "Error parsing argument");
         return NULL;
     }
 
-    int gai_error = 0;
     int mmdb_error = MMDB_SUCCESS;
     MMDB_lookup_result_s result =
-        MMDB_lookup_string(mmdb, ip_address, &gai_error,
-                           &mmdb_error);
-
-    if (0 != gai_error) {
-        PyErr_Format(PyExc_ValueError,
-                     "'%s' does not appear to be an IPv4 or IPv6 address.",
-                     ip_address);
-        return NULL;
-    }
+        MMDB_lookup_sockaddr(mmdb, ip_address, &mmdb_error);
 
     if (MMDB_SUCCESS != mmdb_error) {
         PyObject *exception;
@@ -141,8 +142,11 @@ static PyObject *Reader_get(PyObject *self, PyObject *args)
         } else {
             exception = MaxMindDB_error;
         }
-        PyErr_Format(exception, "Error looking up %s. %s",
-                     ip_address, MMDB_strerror(mmdb_error));
+        char ipstr[INET6_ADDRSTRLEN] = { 0 };
+        if (format_sockaddr(ip_address, ipstr)) {
+            PyErr_Format(exception, "Error looking up %s. %s",
+                         ipstr, MMDB_strerror(mmdb_error));
+        }
         return NULL;
     }
 
@@ -153,9 +157,12 @@ static PyObject *Reader_get(PyObject *self, PyObject *args)
     MMDB_entry_data_list_s *entry_data_list = NULL;
     int status = MMDB_get_entry_data_list(&result.entry, &entry_data_list);
     if (MMDB_SUCCESS != status) {
-        PyErr_Format(MaxMindDB_error,
-                     "Error while looking up data for %s. %s",
-                     ip_address, MMDB_strerror(status));
+        char ipstr[INET6_ADDRSTRLEN] = { 0 };
+        if (format_sockaddr(ip_address, ipstr)) {
+            PyErr_Format(MaxMindDB_error,
+                         "Error while looking up data for %s. %s",
+                         ipstr, MMDB_strerror(status));
+        }
         MMDB_free_entry_data_list(entry_data_list);
         return NULL;
     }
@@ -166,6 +173,113 @@ static PyObject *Reader_get(PyObject *self, PyObject *args)
     return py_obj;
 }
 
+static int ip_converter(PyObject *obj, struct sockaddr_storage *ip_address)
+{
+#if PY_MAJOR_VERSION >= 3
+    if (PyUnicode_Check(obj)) {
+        Py_ssize_t len;
+        const char *ipstr = PyUnicode_AsUTF8AndSize(obj, &len);
+# else
+    if (PyUnicode_Check(obj) || PyString_Check(obj)) {
+        // Although this should work on Python 3, we will hopefully delete
+        // this soon and the Python 3 version is cleaner.
+        const char *ipstr = PyString_AsString(obj);
+        Py_ssize_t len = PyString_Size(obj);
+#endif
+        if (!ipstr) {
+            PyErr_SetString(PyExc_TypeError,
+                            "argument 1 contains an invalid string");
+            return 0;
+        }
+        if (strlen(ipstr) != (size_t)len) {
+            PyErr_SetString(PyExc_TypeError,
+                            "argument 1 contains an embedded null character");
+            return 0;
+        }
+
+        struct addrinfo hints = {
+            .ai_family   = AF_UNSPEC,
+            .ai_flags    = AI_NUMERICHOST,
+            // We set ai_socktype so that we only get one result back
+            .ai_socktype = SOCK_STREAM
+        };
+
+        struct addrinfo *addresses = NULL;
+        int gai_status = getaddrinfo(ipstr, NULL, &hints, &addresses);
+        if (gai_status) {
+            PyErr_Format(PyExc_ValueError,
+                         "'%s' does not appear to be an IPv4 or IPv6 address.",
+                         ipstr);
+            return 0;
+        }
+        if (!addresses) {
+            PyErr_SetString(PyExc_RuntimeError,
+                            "getaddrinfo was successful but failed to set the addrinfo");
+
+        }
+        memcpy(ip_address, addresses->ai_addr, addresses->ai_addrlen);
+        freeaddrinfo(addresses);
+        return 1;
+    }
+    PyObject *packed = PyObject_GetAttrString(obj, "packed");
+    if (!packed) {
+        PyErr_SetString(PyExc_TypeError,
+                        "argument 1 must be a string or ipaddress object");
+        return 0;
+    }
+    Py_ssize_t len;
+    char *bytes;
+    int status = PyBytes_AsStringAndSize(packed, &bytes, &len);
+    if (status == -1) {
+        PyErr_SetString(PyExc_TypeError,
+                        "argument 1 must be a string or ipaddress object");
+        Py_DECREF(packed);
+        return 0;
+    }
+
+    switch (len) {
+    case 16: {
+            ip_address->ss_family = AF_INET6;
+            struct sockaddr_in6 *sin = (struct sockaddr_in6 *)ip_address;
+            memcpy(sin->sin6_addr.s6_addr, bytes, len);
+            Py_DECREF(packed);
+            return 1;
+        }
+    case 4: {
+            ip_address->ss_family = AF_INET;
+            struct sockaddr_in *sin = (struct sockaddr_in *)ip_address;
+            memcpy(&(sin->sin_addr.s_addr), bytes, len);
+            Py_DECREF(packed);
+            return 1;
+        }
+    default:
+        PyErr_SetString(PyExc_ValueError,
+                        "argument 1 returned an unexpected packed length for address");
+        Py_DECREF(packed);
+        return 0;
+    }
+}
+
+static bool format_sockaddr(struct sockaddr *sa, char *dst)
+{
+    char *addr;
+    if (sa->sa_family == AF_INET) {
+        struct sockaddr_in *sin = (struct sockaddr_in *)sa;
+        addr = (char *)&sin->sin_addr;
+    } else {
+        struct sockaddr_in6 *sin = (struct sockaddr_in6 *)sa;
+        addr = (char *)&sin->sin6_addr;
+    }
+
+    if (inet_ntop(sa->sa_family, addr, dst, INET6_ADDRSTRLEN)) {
+        return true;
+    }
+    PyErr_SetString(PyExc_RuntimeError,
+                    "unable to format IP address");
+    return false;
+}
+
+
 static PyObject *Reader_metadata(PyObject *self, PyObject *UNUSED(args))
 {
     Reader_obj *mmdb_obj = (Reader_obj *)self;

maxminddb/compat.py

@@ -24,7 +24,6 @@ def int_from_bytes(b):
 
     string_type = basestring
 
-    string_type_name = 'string'
 else:
 
     def compat_ip_address(address):
@@ -39,5 +38,3 @@ def compat_ip_address(address):
     byte_from_int = lambda x: bytes([x])
 
     string_type = str
-
-    string_type_name = string_type.__name__

maxminddb/reader.py

@@ -15,8 +15,7 @@
 
 import struct
 
-from maxminddb.compat import (byte_from_int, compat_ip_address, string_type,
-                              string_type_name)
+from maxminddb.compat import byte_from_int, compat_ip_address, string_type
 from maxminddb.const import MODE_AUTO, MODE_MMAP, MODE_FILE, MODE_MEMORY, MODE_FD
 from maxminddb.decoder import Decoder
 from maxminddb.errors import InvalidDatabaseError
@@ -105,23 +104,25 @@ def get(self, ip_address):
         Arguments:
         ip_address -- an IP address in the standard string notation
         """
-        if not isinstance(ip_address, string_type):
-            raise TypeError('argument 1 must be %s, not %s' %
-                            (string_type_name, type(ip_address).__name__))
+        if isinstance(ip_address, string_type):
+            address = compat_ip_address(ip_address)
+        else:
+            address = ip_address
 
-        address = compat_ip_address(ip_address)
+        try:
+            packed_address = bytearray(address.packed)
+        except AttributeError:
+            raise TypeError('argument 1 must be a string or ipaddress object')
 
         if address.version == 6 and self._metadata.ip_version == 4:
             raise ValueError(
                 'Error looking up {0}. You attempted to look up '
                 'an IPv6 address in an IPv4-only database.'.format(ip_address))
-        pointer = self._find_address_in_tree(address)
+        pointer = self._find_address_in_tree(packed_address)
 
         return self._resolve_data_pointer(pointer) if pointer else None
 
-    def _find_address_in_tree(self, ip_address):
-        packed = bytearray(ip_address.packed)
-
+    def _find_address_in_tree(self, packed):
         bit_count = len(packed) * 8
         node = self._start_node(bit_count)