commit 7ff9addebe19ae911c30e90296af16838c4896e5 Author: Damian Johnson atagar@torproject.org Date: Sat Jan 27 21:55:31 2018 -0800
Make address type optional for IPv4 and IPv6
Vast majority of time we create addresses they're IPv4 or IPv6. If provided a plaintext address is provided then drop a need for it. --- stem/client/__init__.py | 12 ++++++++++-- test/unit/client/address.py | 18 ++++++++++++------ test/unit/client/cell.py | 8 ++++---- 3 files changed, 26 insertions(+), 12 deletions(-)
diff --git a/stem/client/__init__.py b/stem/client/__init__.py index 58c439e6..fb0f899d 100644 --- a/stem/client/__init__.py +++ b/stem/client/__init__.py @@ -355,7 +355,15 @@ class Address(Field): :var bytes value_bin: encoded address value """
- def __init__(self, addr_type, value): + def __init__(self, value, addr_type = None): + if addr_type is None: + if stem.util.connection.is_valid_ipv4_address(value): + addr_type = AddrType.IPv4 + elif stem.util.connection.is_valid_ipv6_address(value): + addr_type = AddrType.IPv6 + else: + raise ValueError('Address type is required unless an IPv4 or IPv6 address') + self.type, self.type_int = AddrType.get(addr_type)
if self.type == AddrType.IPv4: @@ -408,7 +416,7 @@ class Address(Field):
addr_value, content = split(content, addr_length)
- return Address(addr_type, addr_value), content + return Address(addr_value, addr_type), content
def __hash__(self): return _hash_attr(self, 'type_int', 'value_bin') diff --git a/test/unit/client/address.py b/test/unit/client/address.py index f5dc4a00..c92f7722 100644 --- a/test/unit/client/address.py +++ b/test/unit/client/address.py @@ -31,17 +31,23 @@ class TestAddress(unittest.TestCase): )
for (addr_type, addr_value), expected in test_data: - addr = Address(addr_type, addr_value) + addr = Address(addr_value, addr_type) self.assertEqual(expected.type, addr.type) self.assertEqual(expected.type_int, addr.type_int) self.assertEqual(expected.value, addr.value) self.assertEqual(expected.value_bin, addr.value_bin)
- self.assertRaisesRegexp(ValueError, re.escape("Packed IPv4 addresses should be four bytes, but was: '\x7f\x00'"), Address, 4, '\x7f\x00') - self.assertRaisesRegexp(ValueError, re.escape("Packed IPv6 addresses should be sixteen bytes, but was: '\x7f\x00'"), Address, 6, '\x7f\x00') + # when an IPv4 or IPv6 address the type is optional + + self.assertEqual(AddrType.IPv4, Address('127.0.0.1').type) + self.assertEqual(AddrType.IPv6, Address('2001:0DB8:AC10:FE01::').type) + + self.assertRaisesRegexp(ValueError, re.escape("Packed IPv4 addresses should be four bytes, but was: '\x7f\x00'"), Address, '\x7f\x00', 4) + self.assertRaisesRegexp(ValueError, re.escape("Packed IPv6 addresses should be sixteen bytes, but was: '\x7f\x00'"), Address, '\x7f\x00', 6) + self.assertRaisesRegexp(ValueError, re.escape('Address type is required unless an IPv4 or IPv6 address'), Address, 'nope')
def test_unknown_type(self): - addr = Address(12, 'hello') + addr = Address('hello', 12) self.assertEqual(AddrType.UNKNOWN, addr.type) self.assertEqual(12, addr.type_int) self.assertEqual(None, addr.value) @@ -49,8 +55,8 @@ class TestAddress(unittest.TestCase):
def test_packing(self): test_data = { - '\x04\x04\x7f\x00\x00\x01': Address(AddrType.IPv4, '127.0.0.1'), - '\x06\x10 \x01\r\xb8\x00\x00\x00\x00\x00\x00\xff\x00\x00B\x83)': Address(AddrType.IPv6, '2001:0db8:0000:0000:0000:ff00:0042:8329'), + '\x04\x04\x7f\x00\x00\x01': Address('127.0.0.1'), + '\x06\x10 \x01\r\xb8\x00\x00\x00\x00\x00\x00\xff\x00\x00B\x83)': Address('2001:0db8:0000:0000:0000:ff00:0042:8329'), }
for cell_bytes, address in test_data.items(): diff --git a/test/unit/client/cell.py b/test/unit/client/cell.py index ffead90f..ba28238e 100644 --- a/test/unit/client/cell.py +++ b/test/unit/client/cell.py @@ -6,7 +6,7 @@ import datetime import os import unittest
-from stem.client import ZERO, AddrType, CertType, CloseReason, Address, Certificate +from stem.client import ZERO, CertType, CloseReason, Address, Certificate from test.unit.client import test_data
from stem.client.cell import ( @@ -56,7 +56,7 @@ VERSIONS_CELLS = { }
NETINFO_CELLS = { - '\x00\x00\x08ZZ\xb6\x90\x04\x04\x7f\x00\x00\x01\x01\x04\x04aq\x0f\x02' + ZERO * (FIXED_PAYLOAD_LEN - 17): (datetime.datetime(2018, 1, 14, 1, 46, 56), Address(AddrType.IPv4, '127.0.0.1'), [Address(AddrType.IPv4, '97.113.15.2')]), + '\x00\x00\x08ZZ\xb6\x90\x04\x04\x7f\x00\x00\x01\x01\x04\x04aq\x0f\x02' + ZERO * (FIXED_PAYLOAD_LEN - 17): (datetime.datetime(2018, 1, 14, 1, 46, 56), Address('127.0.0.1'), [Address('97.113.15.2')]), }
VPADDING_CELLS = { @@ -130,8 +130,8 @@ class TestCell(unittest.TestCase): netinfo_cell, content = Cell.unpack(content, 2) self.assertEqual(NetinfoCell, type(netinfo_cell)) self.assertEqual(datetime.datetime(2018, 1, 14, 1, 46, 56), netinfo_cell.timestamp) - self.assertEqual(Address(AddrType.IPv4, '127.0.0.1'), netinfo_cell.receiver_address) - self.assertEqual([Address(AddrType.IPv4, '97.113.15.2')], netinfo_cell.sender_addresses) + self.assertEqual(Address('127.0.0.1'), netinfo_cell.receiver_address) + self.assertEqual([Address('97.113.15.2')], netinfo_cell.sender_addresses)
self.assertEqual('', content) # check that we've consumed all of the bytes