[tor-commits] [stem/master] Add pack instance method to Cell classes

[atagar at torproject.org]

commit eb65f340e0c9cbc429777dc5ce85aede198da138
Author: Damian Johnson <atagar at torproject.org>
Date:   Fri Jan 26 12:46:35 2018 -0800

    Add pack instance method to Cell classes
    
    While a class method is useful, it's also nice to be able to easily pack when
    you have a cell instance.
---
 stem/client/cell.py | 50 ++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 44 insertions(+), 6 deletions(-)

diff --git a/stem/client/cell.py b/stem/client/cell.py
index 4062edcf..37c3d5b9 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -125,6 +125,9 @@ class Cell(object):
     payload, content = split(content, payload_len)
     return cls._unpack(payload, circ_id, link_version), content
 
+  def pack(self, link_version):
+    raise NotImplementedError('Unpacking not yet implemented for %s cells' % cls.NAME)
+
   @classmethod
   def _pack(cls, link_version, payload, circ_id = 0):
     """
@@ -223,6 +226,9 @@ class PaddingCell(Cell):
   def __init__(self, payload):
     self.payload = payload
 
+  def pack(self, link_version):
+    return PaddingCell.pack(link_version, self.payload)
+
   @classmethod
   def pack(cls, link_version, payload = None):
     """
@@ -290,6 +296,9 @@ class RelayCell(CircuitCell):
     elif stream_id and self.command in stem.client.STREAM_ID_DISALLOWED:
       raise ValueError('%s relay cells concern the circuit itself and cannot have a stream id' % self.command)
 
+  def pack(self, link_version):
+    return RelayCell.pack(link_version, self.circ_id, self.command_int, self.data, self.digest, self.stream_id)
+
   @classmethod
   def pack(cls, link_version, circ_id, command, data, digest, stream_id = 0):
     """
@@ -348,6 +357,9 @@ class DestroyCell(CircuitCell):
     super(DestroyCell, self).__init__(circ_id)
     self.reason, self.reason_int = stem.client.CloseReason.get(reason)
 
+  def pack(self, link_version):
+    return DestroyCell.pack(link_version, self.circ_id, self.reason_int)
+
   @classmethod
   def pack(cls, link_version, circ_id, reason = stem.client.CloseReason.NONE):
     """
@@ -389,10 +401,18 @@ class CreateFastCell(CircuitCell):
   VALUE = 5
   IS_FIXED_SIZE = True
 
-  def __init__(self, circ_id, key_material):
+  def __init__(self, circ_id = None, key_material = None):
+    if not key_material:
+      key_material = os.urandom(HASH_LEN)
+    elif len(key_material) != HASH_LEN:
+      raise ValueError('Key material should be %i bytes, but was %i' % (HASH_LEN, len(key_material)))
+
     super(CreateFastCell, self).__init__(circ_id)
     self.key_material = key_material
 
+  def pack(self, link_version):
+    return CreateFastCell.pack(link_version, self.circ_id, self.key_material)
+
   @classmethod
   def pack(cls, link_version, circ_id = None, key_material = None):
     """
@@ -405,16 +425,16 @@ class CreateFastCell(CircuitCell):
     :returns: **bytes** with our randomized key material
     """
 
-    if circ_id is None:
+    cell = CreateFastCell(circ_id, key_material)
+    circ_id = cell.circ_id
+
+    if not circ_id:
       # When initiating a circuit the v4 link protocol requires us to set the
       # most significant bit. Otherwise any id will do.
 
       circ_id = 0x80000000 if link_version >= 4 else 0x01
 
-    if key_material and len(key_material) != HASH_LEN:
-      raise ValueError('Key material should be %i bytes, but was %i' % (HASH_LEN, len(key_material)))
-
-    return cls._pack(link_version, key_material if key_material else os.urandom(HASH_LEN), circ_id)
+    return cls._pack(link_version, cell.key_material, circ_id)
 
   @classmethod
   def _unpack(cls, content, circ_id, link_version):
@@ -446,6 +466,9 @@ class CreatedFastCell(CircuitCell):
     self.key_material = key_material
     self.derivative_key = derivative_key
 
+  def pack(self, link_version):
+    return CreatedFastCell.pack(link_version, self.circ_id, delf.derived_key, self.key_material)
+
   @classmethod
   def pack(cls, link_version, circ_id, derivative_key, key_material = None):
     """
@@ -496,6 +519,9 @@ class VersionsCell(Cell):
   def __init__(self, versions):
     self.versions = versions
 
+  def pack(self, link_version):
+    return VersionsCell.pack(self.versions)
+
   @classmethod
   def pack(cls, versions):
     """
@@ -544,6 +570,9 @@ class NetinfoCell(Cell):
     self.receiver_address = receiver_address
     self.sender_addresses = sender_addresses
 
+  def pack(self, link_version):
+    return NetinfoCell.pack(link_version, self.receiver_address, self.sender_address, self.timestamp)
+
   @classmethod
   def pack(cls, link_version, receiver_address, sender_addresses, timestamp = None):
     """
@@ -629,6 +658,9 @@ class VPaddingCell(Cell):
   def __init__(self, payload):
     self.payload = payload
 
+  def pack(self, link_version):
+    return VPaddingCell.pack(link_version, payload = self.payload)
+
   @classmethod
   def pack(cls, link_version, size = None, payload = None):
     """
@@ -674,6 +706,9 @@ class CertsCell(Cell):
   def __init__(self, certs):
     self.certificates = certs
 
+  def pack(self, link_version):
+    return CertsCell.pack(link_version, self.certificates)
+
   @classmethod
   def pack(cls, link_version, certs):
     """
@@ -722,6 +757,9 @@ class AuthChallengeCell(Cell):
     self.challenge = challenge
     self.methods = methods
 
+  def pack(self, link_version):
+    return AuthChallengeCell.pack(link_version, self.methods, self.challenge)
+
   @classmethod
   def pack(cls, link_version, methods, challenge = None):
     """