Skip to content

Commit 2687c23

Browse files
oschwaldoschwald
committed
Add get_with_prefix_len for pure Python reader
1 parent 8cd9d62 commit 2687c23

File tree

2 files changed

+129
-9
lines changed

2 files changed

+129
-9
lines changed

‎maxminddb/reader.py‎

Lines changed: 24 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,16 @@ def get(self, ip_address):
101101
"""Return the record for the ip_address in the MaxMind DB
102102
103103
104+
Arguments:
105+
ip_address -- an IP address in the standard string notation
106+
"""
107+
(record, _) = self.get_with_prefix_len(ip_address)
108+
return record
109+
110+
def get_with_prefix_len(self, ip_address):
111+
"""Return a tuple with the record and the associated prefix length
112+
113+
104114
Arguments:
105115
ip_address -- an IP address in the standard string notation
106116
"""
@@ -118,24 +128,29 @@ def get(self, ip_address):
118128
raise ValueError(
119129
'Error looking up {0}. You attempted to look up '
120130
'an IPv6 address in an IPv4-only database.'.format(ip_address))
121-
pointer = self._find_address_in_tree(packed_address)
122131

123-
return self._resolve_data_pointer(pointer) if pointer else None
132+
(pointer, prefix_len) = self._find_address_in_tree(packed_address)
133+
134+
if pointer:
135+
return self._resolve_data_pointer(pointer), prefix_len
136+
return None, prefix_len
124137

125138
def _find_address_in_tree(self, packed):
126139
bit_count = len(packed) * 8
127140
node = self._start_node(bit_count)
141+
node_count = self._metadata.node_count
128142

129-
for i in range(bit_count):
130-
if node >= self._metadata.node_count:
131-
break
143+
i = 0
144+
while i < bit_count and node < node_count:
132145
bit = 1 & (packed[i >> 3] >> 7 - (i % 8))
133146
node = self._read_node(node, bit)
134-
if node == self._metadata.node_count:
147+
i = i + 1
148+
149+
if node == node_count:
135150
# Record is empty
136-
return 0
137-
if node > self._metadata.node_count:
138-
return node
151+
return 0, i
152+
if node > node_count:
153+
return node, i
139154

140155
raise InvalidDatabaseError('Invalid node in search tree')
141156

‎tests/reader_test.py‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,111 @@ def test_reader(self):
7272
self._check_ip_v6(reader, file_name)
7373
reader.close()
7474

75+
def test_get_with_prefix_len(self):
76+
decoder_record = {
77+
"array": [1, 2, 3],
78+
"boolean": True,
79+
"bytes": b'\x00\x00\x00*',
80+
"double": 42.123456,
81+
"float": 1.100000023841858,
82+
"int32": -268435456,
83+
"map": {
84+
"mapX": {
85+
"arrayX": [7, 8, 9],
86+
"utf8_stringX": "hello",
87+
},
88+
},
89+
"uint128": 1329227995784915872903807060280344576,
90+
"uint16": 0x64,
91+
"uint32": 0x10000000,
92+
"uint64": 0x1000000000000000,
93+
"utf8_string": "unicode! ☯ - ♫",
94+
}
95+
96+
tests = [{
97+
'ip': '1.1.1.1',
98+
'file_name': 'MaxMind-DB-test-ipv6-32.mmdb',
99+
'expected_prefix_len': 8,
100+
'expected_record': None,
101+
}, {
102+
'ip': '::1:ffff:ffff',
103+
'file_name': 'MaxMind-DB-test-ipv6-24.mmdb',
104+
'expected_prefix_len': 128,
105+
'expected_record': {
106+
"ip": "::1:ffff:ffff"
107+
},
108+
}, {
109+
'ip': '::2:0:1',
110+
'file_name': 'MaxMind-DB-test-ipv6-24.mmdb',
111+
'expected_prefix_len': 122,
112+
'expected_record': {
113+
"ip": "::2:0:0"
114+
},
115+
}, {
116+
'ip': '1.1.1.1',
117+
'file_name': 'MaxMind-DB-test-ipv4-24.mmdb',
118+
'expected_prefix_len': 32,
119+
'expected_record': {
120+
"ip": "1.1.1.1"
121+
},
122+
}, {
123+
'ip': '1.1.1.3',
124+
'file_name': 'MaxMind-DB-test-ipv4-24.mmdb',
125+
'expected_prefix_len': 31,
126+
'expected_record': {
127+
"ip": "1.1.1.2"
128+
},
129+
}, {
130+
'ip': '1.1.1.3',
131+
'file_name': 'MaxMind-DB-test-decoder.mmdb',
132+
'expected_prefix_len': 24,
133+
'expected_record': decoder_record,
134+
}, {
135+
'ip': '::ffff:1.1.1.128',
136+
'file_name': 'MaxMind-DB-test-decoder.mmdb',
137+
'expected_prefix_len': 120,
138+
'expected_record': decoder_record,
139+
}, {
140+
'ip': '::1.1.1.128',
141+
'file_name': 'MaxMind-DB-test-decoder.mmdb',
142+
'expected_prefix_len': 120,
143+
'expected_record': decoder_record,
144+
}, {
145+
'ip': '200.0.2.1',
146+
'file_name': 'MaxMind-DB-no-ipv4-search-tree.mmdb',
147+
'expected_prefix_len': 0,
148+
'expected_record': "::0/64",
149+
}, {
150+
'ip': '::200.0.2.1',
151+
'file_name': 'MaxMind-DB-no-ipv4-search-tree.mmdb',
152+
'expected_prefix_len': 64,
153+
'expected_record': "::0/64",
154+
}, {
155+
'ip': '0:0:0:0:ffff:ffff:ffff:ffff',
156+
'file_name': 'MaxMind-DB-no-ipv4-search-tree.mmdb',
157+
'expected_prefix_len': 64,
158+
'expected_record': "::0/64",
159+
}, {
160+
'ip': 'ef00::',
161+
'file_name': 'MaxMind-DB-no-ipv4-search-tree.mmdb',
162+
'expected_prefix_len': 1,
163+
'expected_record': None,
164+
}]
165+
166+
for test in tests:
167+
reader = open_database('tests/data/test-data/' + test['file_name'],
168+
self.mode)
169+
(record, prefix_len) = reader.get_with_prefix_len(test['ip'])
170+
171+
self.assertEqual(
172+
prefix_len, test['expected_prefix_len'],
173+
'expected prefix_len of {} for {} in {} but got {}'.format(
174+
test['expected_prefix_len'], test['ip'], test['file_name'],
175+
prefix_len))
176+
self.assertEqual(
177+
record, test['expected_record'], 'expected_record for ' +
178+
test['ip'] + ' in ' + test['file_name'])
179+
75180
def test_decoder(self):
76181
reader = open_database(
77182
'tests/data/test-data/MaxMind-DB-test-decoder.mmdb', self.mode)

0 commit comments

Comments
 (0)