[tor-commits] [stem/master] Make address type optional for IPv4 and IPv6

[atagar at torproject.org]

commit 7ff9addebe19ae911c30e90296af16838c4896e5
Author: Damian Johnson <atagar at 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