mmdb_plot.py

#!/usr/bin/python3
from ipaddress import ip_address
import mmap
from sys import argv
import sys

DATA_SECTION_SEP = b"\x00" * 16
META_SECTION_SEP = b"\xab\xcd\xefMaxMind.com"

if __name__ == "__main__":
    with open(argv[1], "rb") as f:
        mm = mmap.mmap(f.fileno(), 0, prot=mmap.PROT_READ)

        data_section_start = mm.find(DATA_SECTION_SEP) + len(DATA_SECTION_SEP)
        meta_section_start = mm.find(META_SECTION_SEP) + len(META_SECTION_SEP)

        binary_tree_size = data_section_start - len(DATA_SECTION_SEP)
        data_section_size = (
            meta_section_start - data_section_start - len(META_SECTION_SEP)
        )

        # Guess node size
        # TODO: Get node size from metadata instead
        # TODO: Handle other node sizes
        if binary_tree_size % 8 == 0:
            node_size = 8
        else:
            raise RuntimeError("Invalid node size")

        node_count = binary_tree_size // node_size
        print(f"node_size={node_size}", file=sys.stderr)
        print(f"node_count={node_count}", file=sys.stderr)

        nodes = {}
        edges = []
        todo = [(0, '')]

        while todo:
            node, bitstring = todo.pop()
            nodes[node] = bitstring

            mm.seek(node * node_size)
            left, right = int.from_bytes(mm.read(node_size // 2)), int.from_bytes(mm.read(node_size // 2))

            # Record value is a node number.
            if left < node_count:
                edges.append((node, left, 0))
                todo.append((left, bitstring + '0'))
            if right < node_count:
                edges.append((node, right, 1))
                todo.append((right, bitstring + '1'))

            # Record value is a pointer in the data section
            if left > node_count:
                edges.append((node, f"data:{left}", 0))
            if right > node_count:
                edges.append((node, f"data:{right}", 1))

        # Guess IP version
        # TODO: Get IP version from metadata instead
        if len(max(nodes.values(), key=len)) < 32:
            tree_height = 32
        else:
            tree_height = 128

        print(f"tree_height={tree_height}", file=sys.stderr)

        print("digraph {")
        for node, bitstring in nodes.items():
            if bitstring == '':
                print(f"{node} [label = \"{node}: root\"];")
            else:
                prefix_int = int(bitstring.ljust(tree_height, '0'), 2)
                prefix = ip_address(prefix_int)
                prefixlen = len(bitstring)
                print(f"{node} [label = \"{node}: {prefix}/{prefixlen}\"];")

        for source, target, label in edges:
            print(f"\"{source}\" -> \"{target}\" [color = \"{'green' if label else 'red'}\", label = \"{label}\"];")
        print("}")