def __init__(self, database):
"""Reader for the MaxMind DB file format
Arguments:
database -- A path to a valid MaxMind DB file such as a GeoIP2
database file.
"""
with open(database, 'rb') as db_file:
self._buffer = mmap.mmap(
db_file.fileno(), 0, access=mmap.ACCESS_READ)
metadata_start = self._buffer.rfind(self._METADATA_START_MARKER,
self._buffer.size() - 128 * 1024)
if metadata_start == -1:
raise InvalidDatabaseError('Error opening database file ({0}). '
'Is this a valid MaxMind DB file?'
''.format(database))
metadata_start += len(self._METADATA_START_MARKER)
metadata_decoder = Decoder(self._buffer, metadata_start)
(metadata, _) = metadata_decoder.decode(metadata_start)
self._metadata = Metadata(**metadata) # pylint: disable=star-args
self._decoder = Decoder(self._buffer, self._metadata.search_tree_size
+ self._DATA_SECTION_SEPARATOR_SIZE)
def __init__(self, database, mode=MODE_AUTO):
"""Reader for the MaxMind DB file format
Arguments:
database -- A path to a valid MaxMind DB file such as a GeoIP2 database
file, or a file descriptor in the case of MODE_FD.
mode -- mode to open the database with. Valid mode are:
* MODE_MMAP - read from memory map.
* MODE_FILE - read database as standard file.
* MODE_MEMORY - load database into memory.
* MODE_AUTO - tries MODE_MMAP and then MODE_FILE. Default.
* MODE_FD - the param passed via database is a file descriptor, not
a path. This mode implies MODE_MEMORY.
"""
if (mode == MODE_AUTO and mmap) or mode == MODE_MMAP:
with open(database, 'rb') as db_file:
self._buffer = mmap.mmap(db_file.fileno(),
0,
access=mmap.ACCESS_READ)
self._buffer_size = self._buffer.size()
filename = database
elif mode in (MODE_AUTO, MODE_FILE):
self._buffer = FileBuffer(database)
self._buffer_size = self._buffer.size()
filename = database
elif mode == MODE_MEMORY:
with open(database, 'rb') as db_file:
self._buffer = db_file.read()
self._buffer_size = len(self._buffer)
filename = database
elif mode == MODE_FD:
self._buffer = database.read()
self._buffer_size = len(self._buffer)
filename = database.name
else:
raise ValueError(
'Unsupported open mode ({0}). Only MODE_AUTO, MODE_FILE, '
'MODE_MEMORY and MODE_FD are supported by the pure Python '
'Reader'.format(mode))
metadata_start = self._buffer.rfind(
self._METADATA_START_MARKER, max(0,
self._buffer_size - 128 * 1024))
if metadata_start == -1:
self.close()
raise InvalidDatabaseError('Error opening database file ({0}). '
'Is this a valid MaxMind DB file?'
''.format(filename))
metadata_start += len(self._METADATA_START_MARKER)
metadata_decoder = Decoder(self._buffer, metadata_start)
(metadata, _) = metadata_decoder.decode(metadata_start)
self._metadata = Metadata(**metadata) # pylint: disable=bad-option-value
self._decoder = Decoder(
self._buffer, self._metadata.search_tree_size +
self._DATA_SECTION_SEPARATOR_SIZE)
self.closed = False
def __init__(self, database, mode=MODE_AUTO):
"""Reader for the MaxMind DB file format
Arguments:
database -- A path to a valid MaxMind DB file such as a GeoIP2 database
file, or a file descriptor in the case of MODE_FD.
mode -- mode to open the database with. Valid mode are:
* MODE_MMAP - read from memory map.
* MODE_FILE - read database as standard file.
* MODE_MEMORY - load database into memory.
* MODE_AUTO - tries MODE_MMAP and then MODE_FILE. Default.
* MODE_FD - the param passed via database is a file descriptor, not
a path. This mode implies MODE_MEMORY.
"""
if (mode == MODE_AUTO and mmap) or mode == MODE_MMAP:
with open(database, 'rb') as db_file:
self._buffer = mmap.mmap(db_file.fileno(),
0,
access=mmap.ACCESS_READ)
self._buffer_size = self._buffer.size()
filename = database
elif mode in (MODE_AUTO, MODE_FILE):
self._buffer = FileBuffer(database)
self._buffer_size = self._buffer.size()
filename = database
elif mode == MODE_MEMORY:
with open(database, 'rb') as db_file:
self._buffer = db_file.read()
self._buffer_size = len(self._buffer)
filename = database
elif mode == MODE_FD:
self._buffer = database.read()
self._buffer_size = len(self._buffer)
filename = database.name
else:
raise ValueError(
'Unsupported open mode ({0}). Only MODE_AUTO, MODE_FILE, '
'MODE_MEMORY and MODE_FD are supported by the pure Python '
'Reader'.format(mode))
metadata_start = self._buffer.rfind(
self._METADATA_START_MARKER, max(0,
self._buffer_size - 128 * 1024))
if metadata_start == -1:
self.close()
raise InvalidDatabaseError('Error opening database file ({0}). '
'Is this a valid MaxMind DB file?'
''.format(filename))
metadata_start += len(self._METADATA_START_MARKER)
metadata_decoder = Decoder(self._buffer, metadata_start)
(metadata, _) = metadata_decoder.decode(metadata_start)
self._metadata = Metadata(**metadata) # pylint: disable=bad-option-value
self._decoder = Decoder(
self._buffer, self._metadata.search_tree_size +
self._DATA_SECTION_SEPARATOR_SIZE)
self.closed = False
def check_decoding(self, type, input, expected, name=None):
name = name or expected
db = mmap.mmap(-1, len(input))
db.write(input)
decoder = Decoder(db, pointer_test=True)
(actual, _,) = decoder.decode(0)
if type in ('float', 'double'):
self.assertAlmostEqual(expected, actual, places=3, msg=type)
else:
self.assertEqual(expected, actual, type)
def test_real_pointers(self):
with open('tests/data/test-data/maps-with-pointers.raw',
'r+b') as db_file:
mm = mmap.mmap(db_file.fileno(), 0)
decoder = Decoder(mm, 0)
self.assertEqual(({
'long_key': 'long_value1'
}, 22), decoder.decode(0))
self.assertEqual(({
'long_key': 'long_value2'
}, 37), decoder.decode(22))
self.assertEqual(({
'long_key2': 'long_value1'
}, 50), decoder.decode(37))
self.assertEqual(({
'long_key2': 'long_value2'
}, 55), decoder.decode(50))
self.assertEqual(({
'long_key': 'long_value1'
}, 57), decoder.decode(55))
self.assertEqual(({
'long_key2': 'long_value2'
}, 59), decoder.decode(57))
mm.close()
def __init__(self, filename):
with open( filename, 'rb' ) as input:
self.data = input.read()
try:
pos = self.data.rindex( MMDB.METADATA_BEGIN_MARKER )
pos += len(MMDB.METADATA_BEGIN_MARKER)
except:
raise Exception( "Could not find metadata marker, invalid database." )
metadata_decoder = Decoder(self.data, pos)
self.metadata = metadata_decoder.decode(pos)[0]
self.metadata = Metadata( **self.metadata )
self.decoder = Decoder(self.data, self.metadata.search_tree_size + MMDB.DATA_SECTION_SEPARATOR_SIZE)
def test_real_pointers(self):
with open("tests/data/test-data/maps-with-pointers.raw",
"r+b") as db_file:
mm = mmap.mmap(db_file.fileno(), 0)
decoder = Decoder(mm, 0)
self.assertEqual(({
"long_key": "long_value1"
}, 22), decoder.decode(0))
self.assertEqual(({
"long_key": "long_value2"
}, 37), decoder.decode(22))
self.assertEqual(({
"long_key2": "long_value1"
}, 50), decoder.decode(37))
self.assertEqual(({
"long_key2": "long_value2"
}, 55), decoder.decode(50))
self.assertEqual(({
"long_key": "long_value1"
}, 57), decoder.decode(55))
self.assertEqual(({
"long_key2": "long_value2"
}, 59), decoder.decode(57))
mm.close()
def test_real_pointers(self):
with open('tests/data/test-data/maps-with-pointers.raw', 'r+b') as db_file:
mm = mmap.mmap(db_file.fileno(), 0)
decoder = Decoder(mm, 0)
self.assertEqual(({'long_key': 'long_value1'}, 22),
decoder.decode(0))
self.assertEqual(({'long_key': 'long_value2'}, 37),
decoder.decode(22))
self.assertEqual(({'long_key2': 'long_value1'}, 50),
decoder.decode(37))
self.assertEqual(({'long_key2': 'long_value2'}, 55),
decoder.decode(50))
self.assertEqual(({'long_key': 'long_value1'}, 57),
decoder.decode(55))
self.assertEqual(({'long_key2': 'long_value2'}, 59),
decoder.decode(57))
mm.close()
class Reader(object):
"""
Instances of this class provide a reader for the MaxMind DB format. IP
addresses can be looked up using the ``get`` method.
"""
_DATA_SECTION_SEPARATOR_SIZE = 16
_METADATA_START_MARKER = b"\xAB\xCD\xEFMaxMind.com"
_ipv4_start = None
def __init__(self, database, mode=MODE_AUTO):
"""Reader for the MaxMind DB file format
Arguments:
database -- A path to a valid MaxMind DB file such as a GeoIP2
database file.
mode -- mode to open the database with. Valid mode are:
* MODE_MMAP - read from memory map.
* MODE_FILE - read database as standard file.
* MODE_MEMORY - load database into memory.
* MODE_AUTO - tries MODE_MMAP and then MODE_FILE. Default.
"""
if (mode == MODE_AUTO and mmap) or mode == MODE_MMAP:
with open(database, 'rb') as db_file:
self._buffer = mmap.mmap(
db_file.fileno(), 0, access=mmap.ACCESS_READ)
self._buffer_size = self._buffer.size()
elif mode in (MODE_AUTO, MODE_FILE):
self._buffer = FileBuffer(database)
self._buffer_size = self._buffer.size()
elif mode == MODE_MEMORY:
with open(database, 'rb') as db_file:
self._buffer = db_file.read()
self._buffer_size = len(self._buffer)
else:
raise ValueError('Unsupported open mode ({0}). Only MODE_AUTO, '
' MODE_FILE, and MODE_MEMORY are support by the pure Python '
'Reader'.format(mode))
metadata_start = self._buffer.rfind(self._METADATA_START_MARKER,
max(0, self._buffer_size
- 128 * 1024))
if metadata_start == -1:
self.close()
raise InvalidDatabaseError('Error opening database file ({0}). '
'Is this a valid MaxMind DB file?'
''.format(database))
metadata_start += len(self._METADATA_START_MARKER)
metadata_decoder = Decoder(self._buffer, metadata_start)
(metadata, _) = metadata_decoder.decode(metadata_start)
self._metadata = Metadata(
**metadata) # pylint: disable=bad-option-value
self._decoder = Decoder(self._buffer, self._metadata.search_tree_size
+ self._DATA_SECTION_SEPARATOR_SIZE)
def metadata(self):
"""Return the metadata associated with the MaxMind DB file"""
return self._metadata
def get(self, ip_address):
"""Return the record for the ip_address in the MaxMind DB
Arguments:
ip_address -- an IP address in the standard string notation
"""
address = ipaddress.ip_address(ip_address)
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)
return self._resolve_data_pointer(pointer) if pointer else None
def _find_address_in_tree(self, ip_address):
packed = ip_address.packed
bit_count = len(packed) * 8
node = self._start_node(bit_count)
for i in range(bit_count):
if node >= self._metadata.node_count:
break
bit = 1 & (int_from_byte(packed[i >> 3]) >> 7 - (i % 8))
node = self._read_node(node, bit)
if node == self._metadata.node_count:
# Record is empty
return 0
elif node > self._metadata.node_count:
return node
raise InvalidDatabaseError('Invalid node in search tree')
def _start_node(self, length):
if self._metadata.ip_version != 6 or length == 128:
return 0
# We are looking up an IPv4 address in an IPv6 tree. Skip over the
# first 96 nodes.
if self._ipv4_start:
return self._ipv4_start
node = 0
for _ in range(96):
if node >= self._metadata.node_count:
break
node = self._read_node(node, 0)
self._ipv4_start = node
return node
def _read_node(self, node_number, index):
base_offset = node_number * self._metadata.node_byte_size
record_size = self._metadata.record_size
if record_size == 24:
offset = base_offset + index * 3
node_bytes = b'\x00' + self._buffer[offset:offset + 3]
elif record_size == 28:
(middle,) = struct.unpack(
b'!B', self._buffer[base_offset + 3:base_offset + 4])
if index:
middle &= 0x0F
else:
middle = (0xF0 & middle) >> 4
offset = base_offset + index * 4
node_bytes = byte_from_int(
middle) + self._buffer[offset:offset + 3]
elif record_size == 32:
offset = base_offset + index * 4
node_bytes = self._buffer[offset:offset + 4]
else:
raise InvalidDatabaseError(
'Unknown record size: {0}'.format(record_size))
return struct.unpack(b'!I', node_bytes)[0]
def _resolve_data_pointer(self, pointer):
resolved = pointer - self._metadata.node_count + \
self._metadata.search_tree_size
if resolved > self._buffer_size:
raise InvalidDatabaseError(
"The MaxMind DB file's search tree is corrupt")
(data, _) = self._decoder.decode(resolved)
return data
def close(self):
"""Closes the MaxMind DB file and returns the resources to the system"""
# pylint: disable=unidiomatic-typecheck
if type(self._buffer) not in (str, bytes):
self._buffer.close()
class Reader:
"""
Instances of this class provide a reader for the MaxMind DB format. IP
addresses can be looked up using the ``get`` method.
"""
_DATA_SECTION_SEPARATOR_SIZE = 16
_METADATA_START_MARKER = b"\xAB\xCD\xEFMaxMind.com"
_buffer: Union[bytes, FileBuffer, "mmap.mmap"]
_ipv4_start: Optional[int] = None
def __init__(self,
database: Union[AnyStr, int, PathLike, IO],
mode: int = MODE_AUTO) -> None:
"""Reader for the MaxMind DB file format
Arguments:
database -- A path to a valid MaxMind DB file such as a GeoIP2 database
file, or a file descriptor in the case of MODE_FD.
mode -- mode to open the database with. Valid mode are:
* MODE_MMAP - read from memory map.
* MODE_FILE - read database as standard file.
* MODE_MEMORY - load database into memory.
* MODE_AUTO - tries MODE_MMAP and then MODE_FILE. Default.
* MODE_FD - the param passed via database is a file descriptor, not
a path. This mode implies MODE_MEMORY.
"""
filename: Any
if (mode == MODE_AUTO and mmap) or mode == MODE_MMAP:
with open(database, "rb") as db_file: # type: ignore
self._buffer = mmap.mmap(db_file.fileno(),
0,
access=mmap.ACCESS_READ)
self._buffer_size = self._buffer.size()
filename = database
elif mode in (MODE_AUTO, MODE_FILE):
self._buffer = FileBuffer(database) # type: ignore
self._buffer_size = self._buffer.size()
filename = database
elif mode == MODE_MEMORY:
with open(database, "rb") as db_file: # type: ignore
self._buffer = db_file.read()
self._buffer_size = len(self._buffer)
filename = database
elif mode == MODE_FD:
self._buffer = database.read() # type: ignore
self._buffer_size = len(self._buffer) # type: ignore
filename = database.name # type: ignore
else:
raise ValueError(
"Unsupported open mode ({0}). Only MODE_AUTO, MODE_FILE, "
"MODE_MEMORY and MODE_FD are supported by the pure Python "
"Reader".format(mode))
metadata_start = self._buffer.rfind(
self._METADATA_START_MARKER, max(0,
self._buffer_size - 128 * 1024))
if metadata_start == -1:
self.close()
raise InvalidDatabaseError("Error opening database file ({0}). "
"Is this a valid MaxMind DB file?"
"".format(filename))
metadata_start += len(self._METADATA_START_MARKER)
metadata_decoder = Decoder(self._buffer, metadata_start)
(metadata, _) = metadata_decoder.decode(metadata_start)
if not isinstance(metadata, dict):
raise InvalidDatabaseError(
"Error reading metadata in database file ({0}).".format(
filename))
self._metadata = Metadata(**metadata) # pylint: disable=bad-option-value
self._decoder = Decoder(
self._buffer,
self._metadata.search_tree_size +
self._DATA_SECTION_SEPARATOR_SIZE,
)
self.closed = False
def metadata(self) -> "Metadata":
"""Return the metadata associated with the MaxMind DB file"""
return self._metadata
def get(self, ip_address: Union[str, IPv6Address,
IPv4Address]) -> Optional[Record]:
"""Return the record for the ip_address in the MaxMind DB
Arguments:
ip_address -- an IP address in the standard string notation
"""
(record, _) = self.get_with_prefix_len(ip_address)
return record
def get_with_prefix_len(
self, ip_address: Union[str, IPv6Address, IPv4Address]
) -> Tuple[Optional[Record], int]:
"""Return a tuple with the record and the associated prefix length
Arguments:
ip_address -- an IP address in the standard string notation
"""
if isinstance(ip_address, str):
address = ipaddress.ip_address(ip_address)
else:
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, prefix_len) = self._find_address_in_tree(packed_address)
if pointer:
return self._resolve_data_pointer(pointer), prefix_len
return None, prefix_len
def _find_address_in_tree(self, packed: bytearray) -> Tuple[int, int]:
bit_count = len(packed) * 8
node = self._start_node(bit_count)
node_count = self._metadata.node_count
i = 0
while i < bit_count and node < node_count:
bit = 1 & (packed[i >> 3] >> 7 - (i % 8))
node = self._read_node(node, bit)
i = i + 1
if node == node_count:
# Record is empty
return 0, i
if node > node_count:
return node, i
raise InvalidDatabaseError("Invalid node in search tree")
def _start_node(self, length: int) -> int:
if self._metadata.ip_version != 6 or length == 128:
return 0
# We are looking up an IPv4 address in an IPv6 tree. Skip over the
# first 96 nodes.
if self._ipv4_start:
return self._ipv4_start
node = 0
for _ in range(96):
if node >= self._metadata.node_count:
break
node = self._read_node(node, 0)
self._ipv4_start = node
return node
def _read_node(self, node_number: int, index: int) -> int:
base_offset = node_number * self._metadata.node_byte_size
record_size = self._metadata.record_size
if record_size == 24:
offset = base_offset + index * 3
node_bytes = b"\x00" + self._buffer[offset:offset + 3]
elif record_size == 28:
offset = base_offset + 3 * index
node_bytes = bytearray(self._buffer[offset:offset + 4])
if index:
node_bytes[0] = 0x0F & node_bytes[0]
else:
middle = (0xF0 & node_bytes.pop()) >> 4
node_bytes.insert(0, middle)
elif record_size == 32:
offset = base_offset + index * 4
node_bytes = self._buffer[offset:offset + 4]
else:
raise InvalidDatabaseError(
"Unknown record size: {0}".format(record_size))
return struct.unpack(b"!I", node_bytes)[0]
def _resolve_data_pointer(self, pointer: int) -> Record:
resolved = pointer - self._metadata.node_count + self._metadata.search_tree_size
if resolved >= self._buffer_size:
raise InvalidDatabaseError(
"The MaxMind DB file's search tree is corrupt")
(data, _) = self._decoder.decode(resolved)
return data
def close(self) -> None:
"""Closes the MaxMind DB file and returns the resources to the system"""
try:
self._buffer.close() # type: ignore
except AttributeError:
pass
self.closed = True
def __exit__(self, *args) -> None:
self.close()
def __enter__(self) -> "Reader":
if self.closed:
raise ValueError("Attempt to reopen a closed MaxMind DB")
return self
class Reader(object):
"""
Instances of this class provide a reader for the MaxMind DB format. IP
addresses can be looked up using the ``get`` method.
"""
_DATA_SECTION_SEPARATOR_SIZE = 16
_METADATA_START_MARKER = b"\xAB\xCD\xEFMaxMind.com"
_ipv4_start = None
def __init__(self, database, mode=MODE_AUTO):
"""Reader for the MaxMind DB file format
Arguments:
database -- A path to a valid MaxMind DB file such as a GeoIP2 database
file, or a file descriptor in the case of MODE_FD.
mode -- mode to open the database with. Valid mode are:
* MODE_MMAP - read from memory map.
* MODE_FILE - read database as standard file.
* MODE_MEMORY - load database into memory.
* MODE_AUTO - tries MODE_MMAP and then MODE_FILE. Default.
* MODE_FD - the param passed via database is a file descriptor, not
a path. This mode implies MODE_MEMORY.
"""
if (mode == MODE_AUTO and mmap) or mode == MODE_MMAP:
with open(database, 'rb') as db_file:
self._buffer = mmap.mmap(db_file.fileno(),
0,
access=mmap.ACCESS_READ)
self._buffer_size = self._buffer.size()
filename = database
elif mode in (MODE_AUTO, MODE_FILE):
self._buffer = FileBuffer(database)
self._buffer_size = self._buffer.size()
filename = database
elif mode == MODE_MEMORY:
with open(database, 'rb') as db_file:
self._buffer = db_file.read()
self._buffer_size = len(self._buffer)
filename = database
elif mode == MODE_FD:
self._buffer = database.read()
self._buffer_size = len(self._buffer)
filename = database.name
else:
raise ValueError(
'Unsupported open mode ({0}). Only MODE_AUTO, MODE_FILE, '
'MODE_MEMORY and MODE_FD are supported by the pure Python '
'Reader'.format(mode))
metadata_start = self._buffer.rfind(
self._METADATA_START_MARKER, max(0,
self._buffer_size - 128 * 1024))
if metadata_start == -1:
self.close()
raise InvalidDatabaseError('Error opening database file ({0}). '
'Is this a valid MaxMind DB file?'
''.format(filename))
metadata_start += len(self._METADATA_START_MARKER)
metadata_decoder = Decoder(self._buffer, metadata_start)
(metadata, _) = metadata_decoder.decode(metadata_start)
self._metadata = Metadata(**metadata) # pylint: disable=bad-option-value
self._decoder = Decoder(
self._buffer, self._metadata.search_tree_size +
self._DATA_SECTION_SEPARATOR_SIZE)
self.closed = False
def metadata(self):
"""Return the metadata associated with the MaxMind DB file"""
return self._metadata
def get(self, ip_address):
"""Return the record for the ip_address in the MaxMind DB
Arguments:
ip_address -- an IP address in the standard string notation
"""
(record, _) = self.get_with_prefix_len(ip_address)
return record
def get_with_prefix_len(self, ip_address):
"""Return a tuple with the record and the associated prefix length
Arguments:
ip_address -- an IP address in the standard string notation
"""
if isinstance(ip_address, string_type):
address = compat_ip_address(ip_address)
else:
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, prefix_len) = self._find_address_in_tree(packed_address)
if pointer:
return self._resolve_data_pointer(pointer), prefix_len
return None, prefix_len
def _find_address_in_tree(self, packed):
bit_count = len(packed) * 8
node = self._start_node(bit_count)
node_count = self._metadata.node_count
i = 0
while i < bit_count and node < node_count:
bit = 1 & (packed[i >> 3] >> 7 - (i % 8))
node = self._read_node(node, bit)
i = i + 1
if node == node_count:
# Record is empty
return 0, i
if node > node_count:
return node, i
raise InvalidDatabaseError('Invalid node in search tree')
def _start_node(self, length):
if self._metadata.ip_version != 6 or length == 128:
return 0
# We are looking up an IPv4 address in an IPv6 tree. Skip over the
# first 96 nodes.
if self._ipv4_start:
return self._ipv4_start
node = 0
for _ in range(96):
if node >= self._metadata.node_count:
break
node = self._read_node(node, 0)
self._ipv4_start = node
return node
def _read_node(self, node_number, index):
base_offset = node_number * self._metadata.node_byte_size
record_size = self._metadata.record_size
if record_size == 24:
offset = base_offset + index * 3
node_bytes = b'\x00' + self._buffer[offset:offset + 3]
elif record_size == 28:
(middle, ) = struct.unpack(
b'!B', self._buffer[base_offset + 3:base_offset + 4])
if index:
middle &= 0x0F
else:
middle = (0xF0 & middle) >> 4
offset = base_offset + index * 4
node_bytes = byte_from_int(middle) + self._buffer[offset:offset +
3]
elif record_size == 32:
offset = base_offset + index * 4
node_bytes = self._buffer[offset:offset + 4]
else:
raise InvalidDatabaseError(
'Unknown record size: {0}'.format(record_size))
return struct.unpack(b'!I', node_bytes)[0]
def _resolve_data_pointer(self, pointer):
resolved = pointer - self._metadata.node_count + \
self._metadata.search_tree_size
if resolved > self._buffer_size:
raise InvalidDatabaseError(
"The MaxMind DB file's search tree is corrupt")
(data, _) = self._decoder.decode(resolved)
return data
def close(self):
"""Closes the MaxMind DB file and returns the resources to the system"""
# pylint: disable=unidiomatic-typecheck
if type(self._buffer) not in (str, bytes):
self._buffer.close()
self.closed = True
def __exit__(self, *args):
self.close()
def __enter__(self):
if self.closed:
raise ValueError('Attempt to reopen a closed MaxMind DB')
return self
def __init__(self,
database: Union[AnyStr, int, PathLike, IO],
mode: int = MODE_AUTO) -> None:
"""Reader for the MaxMind DB file format
Arguments:
database -- A path to a valid MaxMind DB file such as a GeoIP2 database
file, or a file descriptor in the case of MODE_FD.
mode -- mode to open the database with. Valid mode are:
* MODE_MMAP - read from memory map.
* MODE_FILE - read database as standard file.
* MODE_MEMORY - load database into memory.
* MODE_AUTO - tries MODE_MMAP and then MODE_FILE. Default.
* MODE_FD - the param passed via database is a file descriptor, not
a path. This mode implies MODE_MEMORY.
"""
filename: Any
if (mode == MODE_AUTO and mmap) or mode == MODE_MMAP:
with open(database, "rb") as db_file: # type: ignore
self._buffer = mmap.mmap(db_file.fileno(),
0,
access=mmap.ACCESS_READ)
self._buffer_size = self._buffer.size()
filename = database
elif mode in (MODE_AUTO, MODE_FILE):
self._buffer = FileBuffer(database) # type: ignore
self._buffer_size = self._buffer.size()
filename = database
elif mode == MODE_MEMORY:
with open(database, "rb") as db_file: # type: ignore
self._buffer = db_file.read()
self._buffer_size = len(self._buffer)
filename = database
elif mode == MODE_FD:
self._buffer = database.read() # type: ignore
self._buffer_size = len(self._buffer) # type: ignore
filename = database.name # type: ignore
else:
raise ValueError(
f"Unsupported open mode ({mode}). Only MODE_AUTO, MODE_FILE, "
"MODE_MEMORY and MODE_FD are supported by the pure Python "
"Reader")
metadata_start = self._buffer.rfind(
self._METADATA_START_MARKER, max(0,
self._buffer_size - 128 * 1024))
if metadata_start == -1:
self.close()
raise InvalidDatabaseError(
f"Error opening database file ({filename}). "
"Is this a valid MaxMind DB file?")
metadata_start += len(self._METADATA_START_MARKER)
metadata_decoder = Decoder(self._buffer, metadata_start)
(metadata, _) = metadata_decoder.decode(metadata_start)
if not isinstance(metadata, dict):
raise InvalidDatabaseError(
f"Error reading metadata in database file ({filename}).")
self._metadata = Metadata(**metadata) # pylint: disable=bad-option-value
self._decoder = Decoder(
self._buffer,
self._metadata.search_tree_size +
self._DATA_SECTION_SEPARATOR_SIZE,
)
self.closed = False
class Reader(object):
"""
Instances of this class provide a reader for the MaxMind DB format. IP
addresses can be looked up using the ``get`` method.
"""
_DATA_SECTION_SEPARATOR_SIZE = 16
_METADATA_START_MARKER = b"\xAB\xCD\xEFMaxMind.com"
_ipv4_start = None
def __init__(self, database):
"""Reader for the MaxMind DB file format
Arguments:
database -- A path to a valid MaxMind DB file such as a GeoIP2
database file.
"""
with open(database, 'rb') as db_file:
self._buffer = mmap.mmap(
db_file.fileno(), 0, access=mmap.ACCESS_READ)
metadata_start = self._buffer.rfind(self._METADATA_START_MARKER,
self._buffer.size() - 128 * 1024)
if metadata_start == -1:
raise InvalidDatabaseError('Error opening database file ({0}). '
'Is this a valid MaxMind DB file?'
''.format(database))
metadata_start += len(self._METADATA_START_MARKER)
metadata_decoder = Decoder(self._buffer, metadata_start)
(metadata, _) = metadata_decoder.decode(metadata_start)
self._metadata = Metadata(**metadata) # pylint: disable=star-args
self._decoder = Decoder(self._buffer, self._metadata.search_tree_size
+ self._DATA_SECTION_SEPARATOR_SIZE)
def metadata(self):
"""Return the metadata associated with the MaxMind DB file"""
return self._metadata
def get(self, ip_address):
"""Return the record for the ip_address in the MaxMind DB
Arguments:
ip_address -- an IP address in the standard string notation
"""
address = ipaddress.ip_address(ip_address)
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)
return self._resolve_data_pointer(pointer) if pointer else None
def _find_address_in_tree(self, ip_address):
packed = ip_address.packed
bit_count = len(packed) * 8
node = self._start_node(bit_count)
for i in range(bit_count):
if node >= self._metadata.node_count:
break
bit = 1 & (int_from_byte(packed[i >> 3]) >> 7 - (i % 8))
node = self._read_node(node, bit)
if node == self._metadata.node_count:
# Record is empty
return 0
elif node > self._metadata.node_count:
return node
raise InvalidDatabaseError('Invalid node in search tree')
def _start_node(self, length):
if self._metadata.ip_version != 6 or length == 128:
return 0
# We are looking up an IPv4 address in an IPv6 tree. Skip over the
# first 96 nodes.
if self._ipv4_start:
return self._ipv4_start
node = 0
for _ in range(96):
if node >= self._metadata.node_count:
break
node = self._read_node(node, 0)
self._ipv4_start = node
return node
def _read_node(self, node_number, index):
base_offset = node_number * self._metadata.node_byte_size
record_size = self._metadata.record_size
if record_size == 24:
offset = base_offset + index * 3
node_bytes = b'\x00' + self._buffer[offset:offset + 3]
elif record_size == 28:
(middle,) = struct.unpack(
b'!B', self._buffer[base_offset + 3:base_offset + 4])
if index:
middle &= 0x0F
else:
middle = (0xF0 & middle) >> 4
offset = base_offset + index * 4
node_bytes = byte_from_int(
middle) + self._buffer[offset:offset + 3]
elif record_size == 32:
offset = base_offset + index * 4
node_bytes = self._buffer[offset:offset + 4]
else:
raise InvalidDatabaseError(
'Unknown record size: {0}'.format(record_size))
return struct.unpack(b'!I', node_bytes)[0]
def _resolve_data_pointer(self, pointer):
resolved = pointer - self._metadata.node_count + \
self._metadata.search_tree_size
if resolved > self._buffer.size():
raise InvalidDatabaseError(
"The MaxMind DB file's search tree is corrupt")
(data, _) = self._decoder.decode(resolved)
return data
def close(self):
"""Closes the MaxMind DB file and returns the resources to the system"""
self._buffer.close()
class MMDB:
METADATA_BEGIN_MARKER = "\xAB\xCD\xEFMaxMind.com"
DATA_SECTION_SEPARATOR_SIZE = 16
def __init__(self, filename):
with open( filename, 'rb' ) as input:
self.data = input.read()
try:
pos = self.data.rindex( MMDB.METADATA_BEGIN_MARKER )
pos += len(MMDB.METADATA_BEGIN_MARKER)
except:
raise Exception( "Could not find metadata marker, invalid database." )
metadata_decoder = Decoder(self.data, pos)
self.metadata = metadata_decoder.decode(pos)[0]
self.metadata = Metadata( **self.metadata )
self.decoder = Decoder(self.data, self.metadata.search_tree_size + MMDB.DATA_SECTION_SEPARATOR_SIZE)
def dump(self, callback):
root = 0
ip = 0
depth = 1
max_depth = 32 if self.metadata.ip_version == 4 else 128
self._recurse( root, ip, depth, max_depth, callback )
def _recurse(self, root, ip, depth, max_depth, callback):
children = ( self._read_node(root, 0), \
self._read_node(root, 1) )
for bit, child in enumerate(children):
# We ignore empty branches of the search tree
if child == self.metadata.node_count:
continue
if bit == 1:
ip |= ( 1 << ( max_depth - depth ) )
if child < self.metadata.node_count:
self._recurse( child, ip, depth + 1, max_depth, callback )
else:
callback( ip, depth, self._resolve_data_pointer(child) )
def _resolve_data_pointer(self, pointer):
resolved = pointer - self.metadata.node_count + \
self.metadata.search_tree_size
if resolved > len(self.data):
raise Exception("The MaxMind DB file's search tree is corrupt")
(data, _) = self.decoder.decode(resolved)
return data
def _read_node(self, node_number, index):
base_offset = node_number * self.metadata.node_byte_size
record_size = self.metadata.record_size
if record_size == 24:
offset = base_offset + index * 3
node_bytes = b'\x00' + self.data[offset:offset + 3]
elif record_size == 28:
(middle,) = struct.unpack(b'!B', self.data[base_offset + 3:base_offset + 4])
if index:
middle &= 0x0F
else:
middle = (0xF0 & middle) >> 4
offset = base_offset + index * 4
node_bytes = byte_from_int(middle) + self.data[offset:offset + 3]
elif record_size == 32:
offset = base_offset + index * 4
node_bytes = self.data[offset:offset + 4]
else:
raise Exception('Unknown record size: {0}'.format(record_size))
return struct.unpack(b'!I', node_bytes)[0]
