Skip to content

Commit 1942009

Browse files
author
James William Pye
committed
Use ipaddress.
The socket module is not portable on win32 and ipaddress has preliminary status in core.
1 parent 49e8e10 commit 1942009

3 files changed

Lines changed: 106 additions & 136 deletions

File tree

‎postgresql/test/test_driver.py‎

Lines changed: 49 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -278,43 +278,6 @@
278278
Varbit('010111101111'),
279279
],
280280
),
281-
('inet', [
282-
'255.255.255.255',
283-
'127.0.0.1',
284-
'10.0.0.1',
285-
'0.0.0.0',
286-
'::1',
287-
'ffff' + ':ffff'*7,
288-
'fe80::1',
289-
'fe80::1',
290-
'::', # 0::0
291-
],
292-
),
293-
('cidr', [
294-
'255.255.255.255/32',
295-
'127.0.0.0/8',
296-
'127.1.0.0/16',
297-
'10.0.0.0/32',
298-
'0.0.0.0/0',
299-
'ffff' + ':ffff'*7 + '/128',
300-
'::1/128',
301-
'fe80::1/128',
302-
'fe80::/64',
303-
'fe80::/16',
304-
'::/0',
305-
],
306-
),
307-
('inet[]', [
308-
['127.0.0.1', '::1'],
309-
['10.0.0.1', 'fe80::1'],
310-
],
311-
),
312-
('cidr[]', [
313-
['127.0.0.0/8', '::/0'],
314-
['10.0.0.0/16', 'fe80::/64'],
315-
['10.102.0.0/16', 'fe80::/64'],
316-
],
317-
),
318281
('macaddr[]', [
319282
['00:00:00:00:00:00', 'ff:ff:ff:ff:ff:ff'],
320283
['00:00:00:00:00:01', '00:00:00:00:00:00', 'ff:ff:ff:ff:ff:ff'],
@@ -323,6 +286,55 @@
323286
),
324287
]
325288

289+
has_ipaddress = False
290+
try:
291+
import ipaddress
292+
has_ipaddress = True
293+
except ImportError:
294+
import ipaddr as ipaddress
295+
has_ipaddress = True
296+
297+
if has_ipaddress:
298+
type_samples.extend([
299+
('inet', [
300+
ipaddress.IPv4Address('255.255.255.255'),
301+
ipaddress.IPv4Address('127.0.0.1'),
302+
ipaddress.IPv4Address('10.0.0.1'),
303+
ipaddress.IPv4Address('0.0.0.0'),
304+
ipaddress.IPv6Address('::1'),
305+
ipaddress.IPv6Address('ffff' + ':ffff'*7),
306+
ipaddress.IPv6Address('fe80::1'),
307+
ipaddress.IPv6Address('fe80::1'),
308+
ipaddress.IPv6Address('::'), # 0::0
309+
],
310+
),
311+
('cidr', [
312+
ipaddress.IPv4Network('255.255.255.255/32'),
313+
ipaddress.IPv4Network('127.0.0.0/8'),
314+
ipaddress.IPv4Network('127.1.0.0/16'),
315+
ipaddress.IPv4Network('10.0.0.0/32'),
316+
ipaddress.IPv4Network('0.0.0.0/0'),
317+
ipaddress.IPv6Network('ffff' + ':ffff'*7 + '/128'),
318+
ipaddress.IPv6Network('::1/128'),
319+
ipaddress.IPv6Network('fe80::1/128'),
320+
ipaddress.IPv6Network('fe80::/64'),
321+
ipaddress.IPv6Network('fe80::/16'),
322+
ipaddress.IPv6Network('::/0'),
323+
],
324+
),
325+
('inet[]', [
326+
[ipaddress.IPv4Address('127.0.0.1'), ipaddress.IPv6Address('::1')],
327+
[ipaddress.IPv4Address('10.0.0.1'), ipaddress.IPv6Address('fe80::1')],
328+
],
329+
),
330+
('cidr[]', [
331+
[ipaddress.IPv4Network('127.0.0.0/8'), ipaddress.IPv6Network('::/0')],
332+
[ipaddress.IPv4Network('10.0.0.0/16'), ipaddress.IPv6Network('fe80::/64')],
333+
[ipaddress.IPv4Network('10.102.0.0/16'), ipaddress.IPv6Network('fe80::/64')],
334+
],
335+
),
336+
])
337+
326338
class test_driver(unittest.TestCase):
327339
@pg_tmp
328340
def testInterrupt(self):

‎postgresql/types/io/lib.py‎

Lines changed: 27 additions & 91 deletions
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,7 @@ def timetz64_pack(timetup_tz, ql_pack = ql_pack, mktime64 = mktime64):
209209
def timetz64_unpack(data, ql_unpack = ql_unpack, mktimetuple64 = mktimetuple64):
210210
"""
211211
Given "long long" serialized time data, "ql", unpack it into a pair:
212-
212+
213213
((seconds, microseconds), timezone_offset)
214214
"""
215215
ts, tz = ql_unpack(data)
@@ -264,107 +264,44 @@ def varbit_unpack(data, long_unpack = long_unpack):
264264
Given ``varbit`` data, unpack it into a pair:
265265
266266
(bits, data)
267-
267+
268268
Where bits are the total number of bits in data (bytes).
269269
"""
270270
return long_unpack(data[0:4]), data[4:]
271271

272-
from socket import \
273-
AF_INET, AF_INET6, \
274-
inet_pton, inet_ntop
275-
from socket import error as socket_error
276-
# From PGSQL src/include/utils/inet.h
277-
_PGSQL_AF_INET = AF_INET
278-
_PGSQL_AF_INET6 = AF_INET + 1
279-
280-
def net_pack(inet, len = len):
272+
def net_pack(triple,
273+
# Map PGSQL src/include/utils/inet.h to IP version number.
274+
fmap = {
275+
4: 2,
276+
6: 3,
277+
},
278+
len = len,
279+
):
281280
"""
282-
Given a string inet/cidr, yield the serialized form for transport.
283-
284-
Prepends the ``family``, ``mask`` and implicit ``is_cidr`` fields.
281+
net_pack()
285282
286-
Supports cidr and inet types.
283+
Pack Postgres' inet/cidr data structure.
287284
"""
288-
slash_index = inet.find('/')
289-
mask = None
290-
if slash_index >= 0:
291-
try:
292-
mask = int(inet[slash_index+1:])
293-
except ValueError:
294-
raise ValueError('invalid mask in inet/cidr')
295-
address = inet[:slash_index]
296-
else:
297-
address = inet
298-
if inet.find(':') >= 0:
299-
family = _PGSQL_AF_INET6
300-
posix_family = AF_INET6
301-
max_mask = 128
302-
else:
303-
family = _PGSQL_AF_INET
304-
posix_family = AF_INET
305-
max_mask = 32
306-
#If IPv4 address is short, right pad it so that it is valid
307-
num_bytes = len(address.split('.'))
308-
if num_bytes < 4:
309-
address += (4 - num_bytes) * '.0'
310-
try:
311-
data = inet_pton(posix_family, address)
312-
except socket_error as exc:
313-
raise ValueError(str(exc)) from exc
314-
if mask is not None:
315-
if not (0 <= mask <= max_mask):
316-
raise ValueError('invalid mask in inet/cidr')
317-
# Calculate optional cidr byte - PGSQL ignores this on input
318-
i_address = int.from_bytes(data, byteorder='big', signed=False)
319-
i_mask = ~(2**(max_mask-mask)-1)
320-
i_net = i_address & i_mask
321-
is_cidr = 1 if (i_net == i_address) else 0
322-
else:
323-
is_cidr = 0
324-
mask = max_mask
325-
return bytes((family, mask, is_cidr, len(data))) + data
285+
family, mask, data = triple
286+
l = len(data)
287+
return bytes((fmap[family], mask or 0, 0 if mask is None else 1, l)) + data
326288

327-
def net_unpack(data, cidr=False, len=len):
289+
def net_unpack(data,
290+
# Map IP version number to PGSQL src/include/utils/inet.h.
291+
fmap = {
292+
2: 4,
293+
3: 6,
294+
}
295+
):
328296
"""
329-
Given serialized cidr data, return :
297+
net_unpack()
330298
331-
Python string cidr/network representation
299+
Unpack Postgres' inet/cidr data structure.
332300
"""
333301
family, mask, is_cidr, size = data[:4]
302+
return (fmap[family], mask, data[4:])
334303

335-
if family == _PGSQL_AF_INET:
336-
posix_family = AF_INET
337-
max_mask = 32
338-
proper_size = 4
339-
elif family == _PGSQL_AF_INET6:
340-
posix_family = AF_INET6
341-
max_mask = 128
342-
proper_size = 16
343-
else:
344-
raise ValueError("invalid family parameter")
345-
rd = data[4:]
346-
rd_len = len(rd)
347-
if rd_len != size and rd_len != proper_size:
348-
raise ValueError("invalid size parameter")
349-
try:
350-
address = inet_ntop(posix_family, rd)
351-
except socket_error as exc:
352-
raise ValueError(str(exc)) from exc
353-
if not (0 <= mask <= max_mask):
354-
raise ValueError("invalid mask parameter")
355-
if cidr or (mask and mask != max_mask):
356-
result = address + '/' + str(mask)
357-
else:
358-
result = address
359-
return result
360-
361-
def cidr_unpack(data):
362-
"""
363-
Variant of above for CIDR
364-
"""
365-
return net_unpack(data, cidr=True)
366-
367-
def macaddr_pack(data):
304+
def macaddr_pack(data, bytes = bytes):
368305
"""
369306
Pack a MAC address
370307
@@ -394,8 +331,7 @@ def macaddr_pack(data):
394331
raise ValueError('data string cannot be parsed to bytes')
395332
if len(mac_parts) != 6 and len(mac_parts[-1]) != 2:
396333
raise ValueError('data string cannot be parsed to bytes')
397-
macaddr = bytearray([int(p, 16) for p in mac_parts])
398-
return bytes(macaddr)
334+
return bytes([int(p, 16) for p in mac_parts])
399335

400336
def macaddr_unpack(data):
401337
"""

‎postgresql/types/io/pg_network.py‎

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,36 @@
11
from .. import INETOID, CIDROID, MACADDROID
22
from . import lib
3-
4-
oid_to_io = {
5-
MACADDROID : (lib.macaddr_pack, lib.macaddr_unpack, str),
6-
CIDROID : (lib.net_pack, lib.cidr_unpack, str),
7-
INETOID : (lib.net_pack, lib.net_unpack, str),
8-
}
3+
try:
4+
import ipaddress
5+
except ImportError:
6+
import ipaddr as ipaddress
97

108
oid_to_type = {
119
MACADDROID : str,
12-
CIDROID : str,
13-
INETOID : str,
10+
INETOID: ipaddress._IPAddressBase,
11+
CIDROID: ipaddress._BaseNetwork,
12+
}
13+
14+
def inet_pack(ob, pack = lib.net_pack, Constructor = ipaddress.ip_address):
15+
a = Constructor(ob)
16+
return pack((a.version, None, a.packed))
17+
18+
def cidr_pack(ob, pack = lib.net_pack, Constructor = ipaddress.ip_network):
19+
a = Constructor(ob)
20+
return pack((a.version, a.prefixlen, a.network_address.packed))
21+
22+
def inet_unpack(data, unpack = lib.net_unpack, Constructor = ipaddress.ip_address):
23+
version, mask, data = unpack(data)
24+
return Constructor(data)
25+
26+
def cidr_unpack(data, unpack = lib.net_unpack, Constructor = ipaddress.ip_network):
27+
version, mask, data = unpack(data)
28+
r = Constructor(data)
29+
r._prefixlen = mask
30+
return Constructor(str(r))
31+
32+
oid_to_io = {
33+
MACADDROID : (lib.macaddr_pack, lib.macaddr_unpack, str),
34+
CIDROID : (cidr_pack, cidr_unpack, str),
35+
INETOID : (inet_pack, inet_unpack, str),
1436
}

0 commit comments

Comments
 (0)