[tor-commits] [stem/master] Parse address payloads

[atagar at torproject.org]

commit 954269505c8d0656b69849de8121195c42aeee5b
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Jan 15 09:29:45 2018 -0800

    Parse address payloads
    
    Add support for addresses in Netinfo cells. Presently this only decodes IPv4
    addresses since that's what I have an example for.
---
 stem/client/__init__.py      | 74 ++++++++++++++++++++++++++++++++++++++++++++
 test/settings.cfg            |  1 +
 test/unit/client/__init__.py |  1 +
 test/unit/client/types.py    | 19 ++++++++++++
 4 files changed, 95 insertions(+)

diff --git a/stem/client/__init__.py b/stem/client/__init__.py
index 6909c1a7..931dbab1 100644
--- a/stem/client/__init__.py
+++ b/stem/client/__init__.py
@@ -15,17 +15,51 @@ a wrapper for :class:`~stem.socket.RelaySocket`, much the same way as
     |- pack - encodes content
     |- unpack - decodes content
     +- pop - decodes content with remainder
+
+.. data:: AddrType (enum)
+
+  Form an address takes.
+
+  ===================== ===========
+  AddressType           Description
+  ===================== ===========
+  **HOSTNAME**          relay hostname
+  **IPv4**              IPv4 address
+  **IPv6**              IPv6 address
+  **ERROR_TRANSIENT**   temporarily error retrieving address
+  **ERROR_PERMANENT**   permanent error retrieving address
+  **UNKNOWN**           unrecognized address type
+  ===================== ===========
 """
 
 import collections
 import struct
 
+import stem.util.enum
+
 ZERO = '\x00'
 
 __all__ = [
   'cell',
 ]
 
+AddrType = stem.util.enum.UppercaseEnum(
+  'HOSTNAME',
+  'IPv4',
+  'IPv6',
+  'ERROR_TRANSIENT',
+  'ERROR_PERMANENT',
+  'UNKNOWN',
+)
+
+ADDR_INT = {
+  0: AddrType.HOSTNAME,
+  4: AddrType.IPv4,
+  6: AddrType.IPv6,
+  16: AddrType.ERROR_TRANSIENT,
+  17: AddrType.ERROR_PERMANENT,
+}
+
 
 class Certificate(collections.namedtuple('Certificate', ['type', 'value'])):
   """
@@ -45,6 +79,46 @@ class Certificate(collections.namedtuple('Certificate', ['type', 'value'])):
   """
 
 
+class Address(collections.namedtuple('Address', ['type', 'type_int', 'value', 'value_bin', 'ttl'])):
+  """
+  Relay address.
+
+  :var stem.client.AddrType type: address type
+  :var int type_int: integer value of the address type
+  :var unicode value: address value
+  :var bytes value_bin: encoded address value
+  :var int ttl: seconds the record can be validly cached for
+  """
+
+  @staticmethod
+  def pop(content):
+    if not content:
+      raise ValueError('Payload empty where an address was expected')
+    elif len(content) < 2:
+      raise ValueError('Insuffient data for address headers')
+
+    addr_type_int, content = Size.CHAR.pop(content)
+    addr_type = ADDR_INT.get(addr_type_int, AddrType.UNKNOWN)
+    addr_length, content = Size.CHAR.pop(content)
+
+    if len(content) < addr_length:
+      raise ValueError('Address specified a payload of %i bytes, but only had %i' % (addr_length, len(content)))
+    elif len(content) < addr_length + 4:
+      raise ValueError('Address missing a TTL at its end')
+
+    address_bin, content = content[:addr_length], content[addr_length:]
+    ttl, content = Size.LONG.pop(content)
+
+    # TODO: add support for other address types
+
+    address = None
+
+    if addr_type == AddrType.IPv4 and len(address_bin) == 4:
+      address = '.'.join([str(Size.CHAR.unpack(address_bin[i])) for i in range(4)])
+
+    return Address(addr_type, addr_type_int, address, address_bin, ttl), content
+
+
 class Size(object):
   """
   Unsigned `struct.pack format
diff --git a/test/settings.cfg b/test/settings.cfg
index b4b125a8..d4e1a6c3 100644
--- a/test/settings.cfg
+++ b/test/settings.cfg
@@ -229,6 +229,7 @@ test.unit_tests
 |test.unit.response.authchallenge.TestAuthChallengeResponse
 |test.unit.response.protocolinfo.TestProtocolInfoResponse
 |test.unit.response.mapaddress.TestMapAddressResponse
+|test.unit.client.types.TestClientTypes
 |test.unit.client.cell.TestCell
 |test.unit.connection.authentication.TestAuthenticate
 |test.unit.connection.connect.TestConnect
diff --git a/test/unit/client/__init__.py b/test/unit/client/__init__.py
index f5b82e7e..fdc7a0c6 100644
--- a/test/unit/client/__init__.py
+++ b/test/unit/client/__init__.py
@@ -4,6 +4,7 @@ Unit tests for stem.client.* contents.
 
 __all__ = [
   'cell',
+  'types',
 ]
 
 import os
diff --git a/test/unit/client/types.py b/test/unit/client/types.py
new file mode 100644
index 00000000..fdf4e085
--- /dev/null
+++ b/test/unit/client/types.py
@@ -0,0 +1,19 @@
+"""
+Unit tests for the types in stem.client.
+"""
+
+import unittest
+
+from stem.client import Address
+
+
+class TestClientTypes(unittest.TestCase):
+  def test_address_ipv4(self):
+    addr, content = Address.pop('\x04\x04\x7f\x00\x00\x01\x01\x04\x04aq\x0f\x02\x00\x00\x00\x00')
+    self.assertEqual('q\x0f\x02\x00\x00\x00\x00', content)
+
+    self.assertEqual('IPv4', addr.type)
+    self.assertEqual(4, addr.type_int)
+    self.assertEqual('127.0.0.1', addr.value)
+    self.assertEqual('\x7f\x00\x00\x01', addr.value_bin)
+    self.assertEqual(17040481, addr.ttl)