commit a615cf7bc392337cc4353885ab64530abeae3dcc Author: Damian Johnson atagar@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)