[tor-commits] [stem/master] Replace Pack enum with a Size class

[atagar at torproject.org]

commit a615cf7bc392337cc4353885ab64530abeae3dcc
Author: Damian Johnson <atagar at torproject.org>
Date:   Thu Jan 11 11:56:32 2018 -0800

    Replace Pack enum with a Size class
    
    Swapping our enum with the class so it can both pack and unpack.
---
 stem/client/__init__.py | 72 ++++++++++++++++++++++++++++++++++++++++---------
 stem/client/cell.py     | 11 ++++----
 2 files changed, 65 insertions(+), 18 deletions(-)

diff --git a/stem/client/__init__.py b/stem/client/__init__.py
index 21786835..d0920778 100644
--- a/stem/client/__init__.py
+++ b/stem/client/__init__.py
@@ -28,29 +28,77 @@ a wrapper for :class:`~stem.socket.RelaySocket`, much the same way as
 
 ::
 
-  Relay - Connection with a relay's ORPort.
+  Size - Packable and unpackable field size.
+    |- pack - encodes content
+    |- unpack - decodes content
+    +- pop - decodes content with remainder
 """
 
-from stem.util import enum
+import struct
 
 ZERO = '\x00'
 
-
 __all__ = [
   'cell',
 ]
 
 
-class Relay(object):
+class Size(object):
   """
-  Connection with a `Tor relay's ORPort
-  <https://gitweb.torproject.org/torspec.git/tree/tor-spec.txt>`_.
+  Unsigned `struct.pack format
+  <https://docs.python.org/2/library/struct.html#format-characters>` for
+  network-order fields.
   """
 
+  def __init__(self, name, size, pack_format):
+    self.name = name
+    self.size = size
+    self.format = pack_format
+
+  def pack(self, content):
+    """
+    Encodes bytes into a packed field.
+
+    :param bytes content: content to encode
+
+    :raises: **ValueError** if content isn't of the right size
+    """
+
+    unpacked = struct.pack(self.format, content)
+
+    if self.size != len(unpacked):
+      raise ValueError("'%s' is the wrong size for a %s field" % (unpacked, self.name))
+
+    return unpacked
+
+  def unpack(self, content):
+    """
+    Decodes packed data into bytes.
+
+    :param bytes content: content to encode
+
+    :raises: **ValueError** if packed data isn't of the right size
+    """
+
+    if self.size != len(content):
+      raise ValueError("'%s' is the wrong size for a %s field" % (content, self.name))
+
+    return struct.unpack(self.format, content)[0]
+
+  def pop(self, content):
+    """
+    Decodes the first characters as this data type, providing it and the
+    remainder.
+
+    :param bytes content: content to encode
+
+    :raises: **ValueError** if packed data isn't of the right size
+    """
+
+    return self.unpack(content[:self.size]), content[self.size:]
+
 
-Pack = enum.Enum(
-  ('CHAR', '!B'),       # 1 byte
-  ('SHORT', '!H'),      # 2 bytes
-  ('LONG', '!L'),       # 4 bytes
-  ('LONG_LONG', '!Q'),  # 8 bytes
-)
+setattr(Size, 'CHAR', Size('CHAR', 1, '!B'))
+setattr(Size, 'SHORT', Size('SHORT', 2, '!H'))
+setattr(Size, 'LONG', Size('LONG', 4, '!L'))
+setattr(Size, 'LONG_LONG', Size('LONG_LONG', 8, '!Q'))
diff --git a/stem/client/cell.py b/stem/client/cell.py
index 7c680c4e..42fe87e1 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -38,11 +38,10 @@ Messages communicated over a Tor relay's ORPort.
 
 import collections
 import inspect
-import struct
 import sys
 
 from stem import UNDEFINED
-from stem.client import ZERO, Pack
+from stem.client import ZERO, Size
 
 
 class Cell(collections.namedtuple('Cell', ['name', 'value', 'fixed_size', 'for_circuit'])):
@@ -124,9 +123,9 @@ class Cell(collections.namedtuple('Cell', ['name', 'value', 'fixed_size', 'for_c
     :raise: **ValueError** if cell type invalid or payload is too large
     """
 
-    packed_circ_id = struct.pack(Pack.LONG if link_version > 3 else Pack.SHORT, circ_id)
-    packed_command = struct.pack(Pack.CHAR, cls.VALUE)
-    packed_size = b'' if cls.IS_FIXED_SIZE else struct.pack(Pack.SHORT, len(payload))
+    packed_circ_id = Size.LONG.pack(circ_id) if link_version > 3 else Size.SHORT.pack(circ_id)
+    packed_command = Size.CHAR.pack(cls.VALUE)
+    packed_size = b'' if cls.IS_FIXED_SIZE else Size.SHORT.pack(len(payload))
     cell = b''.join((packed_circ_id, packed_command, packed_size, payload))
 
     # pad fixed sized cells to the required length
@@ -249,7 +248,7 @@ class VersionsCell(Cell):
     # Used for link version negotiation so we don't have that yet. This is fine
     # since VERSION cells avoid most version dependent attributes.
 
-    payload = b''.join([struct.pack(Pack.SHORT, v) for v in versions])
+    payload = b''.join([Size.SHORT.pack(v) for v in versions])
     return cls._pack(3, payload)