[tor-commits] [stem/master] Support Address packing

[atagar at torproject.org]

commit e8490b15d6585a8fb33d6d8d7002c69cf68ba120
Author: Damian Johnson <atagar at torproject.org>
Date:   Sat Jan 20 12:39:45 2018 -0800

    Support Address packing
---
 stem/client/__init__.py     |  8 ++++++++
 test/unit/client/address.py | 18 +++++++++++++++++-
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/stem/client/__init__.py b/stem/client/__init__.py
index 7705e59d..13952a00 100644
--- a/stem/client/__init__.py
+++ b/stem/client/__init__.py
@@ -52,6 +52,7 @@ a wrapper for :class:`~stem.socket.RelaySocket`, much the same way as
 """
 
 import collections
+import io
 import struct
 
 import stem.util.connection
@@ -273,6 +274,13 @@ class Address(Field):
       self.value = None
       self.value_bin = value
 
+  def pack(self):
+    cell = io.BytesIO()
+    cell.write(Size.CHAR.pack(self.type_int))
+    cell.write(Size.CHAR.pack(len(self.value_bin)))
+    cell.write(self.value_bin)
+    return cell.getvalue()
+
   @staticmethod
   def pop(content):
     if not content:
diff --git a/test/unit/client/address.py b/test/unit/client/address.py
index 689b77a0..acda251b 100644
--- a/test/unit/client/address.py
+++ b/test/unit/client/address.py
@@ -33,7 +33,23 @@ class TestAddress(unittest.TestCase):
     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')
 
-  def test_pop_ipv4(self):
+  def test_unknown_type(self):
+    addr = Address(12, 'hello')
+    self.assertEqual(AddrType.UNKNOWN, addr.type)
+    self.assertEqual(12, addr.type_int)
+    self.assertEqual(None, addr.value)
+    self.assertEqual('hello', addr.value_bin)
+
+  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'),
+    }
+
+    for cell_bytes, address in test_data.items():
+      self.assertEqual(cell_bytes, address.pack())
+      self.assertEqual(address, Address.unpack(cell_bytes))
+
     addr, content = Address.pop('\x04\x04\x7f\x00\x00\x01\x01\x04\x04aq\x0f\x02\x00\x00\x00\x00')
     self.assertEqual('\x01\x04\x04aq\x0f\x02\x00\x00\x00\x00', content)