"""Pure-Python reader for the MaxMind DB file format.""" from __future__ import annotations try: import mmap except ImportError: mmap = None # type: ignore[assignment] import contextlib import ipaddress import struct from ipaddress import IPv4Address, IPv6Address from typing import IO, TYPE_CHECKING, Any, AnyStr from maxminddb.const import MODE_AUTO, MODE_FD, MODE_FILE, MODE_MEMORY, MODE_MMAP from maxminddb.decoder import Decoder from maxminddb.errors import InvalidDatabaseError from maxminddb.file import FileBuffer if TYPE_CHECKING: from collections.abc import Iterator from os import PathLike from typing_extensions import Self from maxminddb.types import Record _IPV4_MAX_NUM = 2**32 class Reader: """A pure Python implementation of 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: bytes | FileBuffer | "mmap.mmap" # noqa: UP037 _buffer_size: int closed: bool _decoder: Decoder _metadata: Metadata _ipv4_start: int def __init__( self, database: 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[arg-type] 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[arg-type] self._buffer_size = self._buffer.size() filename = database elif mode == MODE_MEMORY: with open(database, "rb") as db_file: # type: ignore[arg-type] buf = db_file.read() self._buffer = buf self._buffer_size = len(buf) filename = database elif mode == MODE_FD: self._buffer = database.read() # type: ignore[union-attr] self._buffer_size = len(self._buffer) # type: ignore[arg-type] filename = database.name # type: ignore[union-attr] else: msg = ( f"Unsupported open mode ({mode}). Only MODE_AUTO, MODE_FILE, " "MODE_MEMORY and MODE_FD are supported by the pure Python " "Reader" ) raise ValueError( msg, ) metadata_start = self._buffer.rfind( self._METADATA_START_MARKER, max(0, self._buffer_size - 128 * 1024), ) if metadata_start == -1: self.close() msg = ( f"Error opening database file ({filename}). " "Is this a valid MaxMind DB file?" ) raise InvalidDatabaseError( msg, ) 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): msg = f"Error reading metadata in database file ({filename})." raise InvalidDatabaseError( msg, ) self._metadata = Metadata(**metadata) self._decoder = Decoder( self._buffer, self._metadata.search_tree_size + self._DATA_SECTION_SEPARATOR_SIZE, ) self.closed = False ipv4_start = 0 if self._metadata.ip_version == 6: # We store the IPv4 starting node as an optimization for IPv4 lookups # in IPv6 trees. This allows us to skip over the first 96 nodes in # this case. node = 0 for _ in range(96): if node >= self._metadata.node_count: break node = self._read_node(node, 0) ipv4_start = node self._ipv4_start = ipv4_start def metadata(self) -> Metadata: """Return the metadata associated with the MaxMind DB file.""" return self._metadata def get(self, ip_address: str | IPv6Address | IPv4Address) -> Record | None: """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: str | IPv6Address | IPv4Address, ) -> tuple[Record | None, 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 as ex: msg = "argument 1 must be a string or ipaddress object" raise TypeError(msg) from ex if address.version == 6 and self._metadata.ip_version == 4: msg = ( f"Error looking up {ip_address}. You attempted to look up " "an IPv6 address in an IPv4-only database." ) raise ValueError( msg, ) (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 __iter__(self) -> Iterator: return self._generate_children(0, 0, 0) def _generate_children(self, node: int, depth: int, ip_acc: int) -> Iterator: if ip_acc != 0 and node == self._ipv4_start: # Skip nodes aliased to IPv4 return node_count = self._metadata.node_count if node > node_count: bits = 128 if self._metadata.ip_version == 6 else 32 ip_acc <<= bits - depth if ip_acc <= _IPV4_MAX_NUM and bits == 128: depth -= 96 yield ( ipaddress.ip_network((ip_acc, depth)), self._resolve_data_pointer( node, ), ) elif node < node_count: left = self._read_node(node, 0) ip_acc <<= 1 depth += 1 yield from self._generate_children(left, depth, ip_acc) right = self._read_node(node, 1) yield from self._generate_children(right, depth, ip_acc | 1) 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 msg = "Invalid node in search tree" raise InvalidDatabaseError(msg) def _start_node(self, length: int) -> int: if self._metadata.ip_version == 6 and length == 32: return self._ipv4_start return 0 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: msg = f"Unknown record size: {record_size}" raise InvalidDatabaseError(msg) 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: msg = "The MaxMind DB file's search tree is corrupt" raise InvalidDatabaseError(msg) (data, _) = self._decoder.decode(resolved) return data def close(self) -> None: """Close the MaxMind DB file and returns the resources to the system.""" with contextlib.suppress(AttributeError): self._buffer.close() # type: ignore[union-attr] self.closed = True def __exit__(self, *_) -> None: # noqa: ANN002 self.close() def __enter__(self) -> Self: if self.closed: msg = "Attempt to reopen a closed MaxMind DB" raise ValueError(msg) return self class Metadata: """Metadata for the MaxMind DB reader.""" binary_format_major_version: int """ The major version number of the binary format used when creating the database. """ binary_format_minor_version: int """ The minor version number of the binary format used when creating the database. """ build_epoch: int """ The Unix epoch for the build time of the database. """ database_type: str """ A string identifying the database type, e.g., "GeoIP2-City". """ description: dict[str, str] """ A map from locales to text descriptions of the database. """ ip_version: int """ The IP version of the data in a database. A value of "4" means the database only supports IPv4. A database with a value of "6" may support both IPv4 and IPv6 lookups. """ languages: list[str] """ A list of locale codes supported by the database. """ node_count: int """ The number of nodes in the database. """ record_size: int """ The bit size of a record in the search tree. """ def __init__(self, **kwargs) -> None: """Create new Metadata object. kwargs are key/value pairs from spec.""" # Although I could just update __dict__, that is less obvious and it # doesn't work well with static analysis tools and some IDEs self.node_count = kwargs["node_count"] self.record_size = kwargs["record_size"] self.ip_version = kwargs["ip_version"] self.database_type = kwargs["database_type"] self.languages = kwargs["languages"] self.binary_format_major_version = kwargs["binary_format_major_version"] self.binary_format_minor_version = kwargs["binary_format_minor_version"] self.build_epoch = kwargs["build_epoch"] self.description = kwargs["description"] @property def node_byte_size(self) -> int: """The size of a node in bytes. :type: int """ return self.record_size // 4 @property def search_tree_size(self) -> int: """The size of the search tree. :type: int """ return self.node_count * self.node_byte_size def __repr__(self) -> str: args = ", ".join(f"{k}={v!r}" for k, v in self.__dict__.items()) return f"{self.__module__}.{self.__class__.__name__}({args})"