Move cell encryption/decription into the RelayCell class

"Less is more."

Dave's pull request for ticket #27112 refactored our RelayCell into a
BaseRelayCell and RawRelayCell class. Think I see what he's going for,
and this was a step forward in some respects. But honestly added a lot
of code and complexity without functionally doing anything new.

Going with a hybrid of our approaches that integrates the part I think
we agree on. Namely, simplifying our Circuit's send() method by moving
cell encryption/decription down into the RelayCell class.

This is a much, *MUCH* smaller diff than what Dave proposed. Happy to
discuss integrating more of the pull request, but lets start with
the part I think we agree on to give us a common point to work from.

8 files changed, 145 insertions, 56 deletions

diff --git a/stem/client/__init__.py b/stem/client/__init__.py
index 24b6de38..aaa62f5b 100644
--- a/stem/client/__init__.py
+++ b/stem/client/__init__.py
@@ -25,7 +25,6 @@ a wrapper for :class:`~stem.socket.RelaySocket`, much the same way as
     +- close - closes this circuit
 """
 
-import copy
 import hashlib
 import threading
 
@@ -34,7 +33,7 @@ import stem.client.cell
 import stem.socket
 import stem.util.connection
 
-from stem.client.datatype import ZERO, LinkProtocol, Address, Size, KDF, split
+from stem.client.datatype import ZERO, LinkProtocol, Address, KDF, split
 
 __all__ = [
   'cell',
@@ -226,7 +225,7 @@ class Circuit(object):
     """
     Sends a message over the circuit.
 
-    :param stem.client.RelayCommand command: command to be issued
+    :param stem.client.datatype.RelayCommand command: command to be issued
     :param bytes data: message payload
     :param int stream_id: specific stream this concerns
 
@@ -234,54 +233,38 @@ class Circuit(object):
     """
 
     with self.relay._orport_lock:
-      orig_digest = self.forward_digest.copy()
-      orig_key = copy.copy(self.forward_key)
+      # Encrypt and send the cell. Our digest/key only updates if the cell is
+      # successfully sent.
 
-      # Digests and such are computed using the RELAY cell payload. This
-      # doesn't include the initial circuit id and cell type fields.
-      # Circuit ids vary in length depending on the protocol version.
+      cell = stem.client.cell.RelayCell(self.id, command, data, stream_id = stream_id)
+      payload, forward_key, forward_digest = cell.encrypt(self.relay.link_protocol, self.forward_key, self.forward_digest)
+      self.relay._orport.send(payload)
 
-      header_size = self.relay.link_protocol.circ_id_size.size + 1
+      self.forward_digest = forward_digest
+      self.forward_key = forward_key
 
-      try:
-        cell = stem.client.cell.RelayCell(self.id, command, data, 0, stream_id)
-        payload_without_digest = cell.pack(self.relay.link_protocol)[header_size:]
-        self.forward_digest.update(payload_without_digest)
+      # Decrypt relay cells received in response. Again, our digest/key only
+      # updates when handled successfully.
 
-        cell = stem.client.cell.RelayCell(self.id, command, data, self.forward_digest, stream_id)
-        header, payload = split(cell.pack(self.relay.link_protocol), header_size)
-        encrypted_payload = header + self.forward_key.update(payload)
+      reply = self.relay._orport.recv()
+      reply_cells = []
 
-        reply_cells = []
-        self.relay._orport.send(encrypted_payload)
-        reply = self.relay._orport.recv()
+      if len(reply) % self.relay.link_protocol.fixed_cell_length != 0:
+        raise stem.ProtocolError('Circuit response should be a series of RELAY cells, but received an unexpected size for a response: %i' % len(reply))
 
-        # Check that we got the correct number of bytes for a series of RELAY cells
+      while reply:
+        encrypted_cell, reply = split(reply, self.relay.link_protocol.fixed_cell_length)
+        decrypted_cell, backward_key, backward_digest = stem.client.cell.RelayCell.decrypt(self.relay.link_protocol, encrypted_cell, self.backward_key, self.backward_digest)
 
-        relay_cell_size = header_size + stem.client.cell.FIXED_PAYLOAD_LEN
-        relay_cell_cmd = stem.client.cell.RelayCell.VALUE
+        if self.id != decrypted_cell.circ_id:
+          raise stem.ProtocolError('Response should be for circuit id %i, not %i' % (self.id, decrypted_cell.circ_id))
 
-        if len(reply) % relay_cell_size != 0:
-          raise stem.ProtocolError('Circuit response should be a series of RELAY cells, but received an unexpected size for a response: %i' % len(reply))
+        self.backward_digest = backward_digest
+        self.backward_key = backward_key
 
-        while reply:
-          circ_id, reply = self.relay.link_protocol.circ_id_size.pop(reply)
-          command, reply = Size.CHAR.pop(reply)
-          payload, reply = split(reply, stem.client.cell.FIXED_PAYLOAD_LEN)
+        reply_cells.append(decrypted_cell)
 
-          if command != relay_cell_cmd:
-            raise stem.ProtocolError('RELAY cell responses should be %i but was %i' % (relay_cell_cmd, command))
-          elif circ_id != self.id:
-            raise stem.ProtocolError('Response should be for circuit id %i, not %i' % (self.id, circ_id))
-
-          decrypted = self.backward_key.update(payload)
-          reply_cells.append(stem.client.cell.RelayCell._unpack(decrypted, self.id, self.relay.link_protocol))
-
-        return reply_cells
-      except:
-        self.forward_digest = orig_digest
-        self.forward_key = orig_key
-        raise
+      return reply_cells
 
   def close(self):
     with self.relay._orport_lock:
diff --git a/stem/client/cell.py b/stem/client/cell.py
index 12fa994c..91dec143 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -37,6 +37,7 @@ Messages communicated over a Tor relay's ORPort.
     +- pop - decodes cell with remainder
 """
 
+import copy
 import datetime
 import inspect
 import os
@@ -101,9 +102,9 @@ class Cell(object):
     """
     Provides cell attributes by its name.
 
-    :parm str name: cell command to fetch
+    :param str name: cell command to fetch
 
-    :raise: **ValueError** if cell type is invalid
+    :raises: **ValueError** if cell type is invalid
     """
 
     for _, cls in inspect.getmembers(sys.modules[__name__]):
@@ -117,9 +118,9 @@ class Cell(object):
     """
     Provides cell attributes by its value.
 
-    :parm int value: cell value to fetch
+    :param int value: cell value to fetch
 
-    :raise: **ValueError** if cell type is invalid
+    :raises: **ValueError** if cell type is invalid
     """
 
     for _, cls in inspect.getmembers(sys.modules[__name__]):
@@ -199,9 +200,9 @@ class Cell(object):
     :param bytes payload: cell payload
     :param int circ_id: circuit id, if a CircuitCell
 
-    :return: **bytes** with the encoded payload
+    :returns: **bytes** with the encoded payload
 
-    :raise: **ValueError** if cell type invalid or payload makes cell too large
+    :raises: **ValueError** if cell type invalid or payload makes cell too large
     """
 
     if issubclass(cls, CircuitCell):
@@ -324,10 +325,16 @@ class RelayCell(CircuitCell):
   """
   Command concerning a relay circuit.
 
+  Our 'recognized' attribute provides a cheap (but incomplete) check for if our
+  cell payload is encrypted. If non-zero our payload *IS* encrypted, but if
+  zero we're *PROBABLY* fully decrypted. This uncertainty is because encrypted
+  cells have a small chance of coincidently producing zero for this value as
+  well.
+
   :var stem.client.RelayCommand command: command to be issued
   :var int command_int: integer value of our command
   :var bytes data: payload of the cell
-  :var int recognized: zero if cell is decrypted, non-zero otherwise
+  :var int recognized: non-zero if payload is encrypted
   :var int digest: running digest held with the relay
   :var int stream_id: specific stream this concerns
   """
@@ -376,6 +383,96 @@ class RelayCell(CircuitCell):
 
     return RelayCell._pack(link_protocol, bytes(payload), self.unused, self.circ_id)
 
+  @staticmethod
+  def decrypt(link_protocol, content, key, digest):
+    """
+    Decrypts content as a relay cell addressed to us. This provides back a
+    tuple of the form...
+
+    ::
+
+      (cell (RelayCell), new_key (CipherContext), new_digest (HASH))
+
+    :param int link_protocol: link protocol version
+    :param bytes content: cell content to be decrypted
+    :param cryptography.hazmat.primitives.ciphers.CipherContext key:
+      key established with the relay we received this cell from
+    :param HASH digest: running digest held with the relay
+
+    :returns: **tuple** with our decrypted cell and updated key/digest
+
+    :raises: :class:`stem.ProtocolError` if content doesn't belong to a relay
+      cell
+    """
+
+    new_key = copy.copy(key)
+    new_digest = digest.copy()
+
+    if len(content) != link_protocol.fixed_cell_length:
+      raise stem.ProtocolError('RELAY cells should be %i bytes, but received %i' % (link_protocol.fixed_cell_length, len(content)))
+
+    circ_id, content = link_protocol.circ_id_size.pop(content)
+    command, encrypted_payload = Size.CHAR.pop(content)
+
+    if command != RelayCell.VALUE:
+      raise stem.ProtocolError('Cannot decrypt as a RELAY cell. This had command %i instead.' % command)
+
+    payload = new_key.update(encrypted_payload)
+
+    cell = RelayCell._unpack(payload, circ_id, link_protocol)
+
+    # TODO: Implement our decryption digest. It is used to support relaying
+    # within multi-hop circuits. On first glance this should go something
+    # like...
+    #
+    #   # Our updated digest is calculated based on this cell with a blanked
+    #   # digest field.
+    #
+    #   digest_cell = RelayCell(self.circ_id, self.command, self.data, 0, self.stream_id, self.recognized, self.unused)
+    #   new_digest.update(digest_cell.pack(link_protocol))
+    #
+    #   is_encrypted == cell.recognized != 0 or self.digest == new_digest
+    #
+    # ... or something like that. Until we attempt to support relaying this is
+    # both moot and difficult to exercise in order to ensure we get it right.
+
+    return cell, new_key, new_digest
+
+  def encrypt(self, link_protocol, key, digest):
+    """
+    Encrypts our cell content to be sent with the given key. This provides back
+    a tuple of the form...
+
+    ::
+
+      (payload (bytes), new_key (CipherContext), new_digest (HASH))
+
+    :param int link_protocol: link protocol version
+    :param cryptography.hazmat.primitives.ciphers.CipherContext key:
+      key established with the relay we're sending this cell to
+    :param HASH digest: running digest held with the relay
+
+    :returns: **tuple** with our encrypted payload and updated key/digest
+    """
+
+    new_key = copy.copy(key)
+    new_digest = digest.copy()
+
+    # Digests are computed from our payload, not including our header's circuit
+    # id (2 or 4 bytes) and command (1 byte).
+
+    header_size = link_protocol.circ_id_size.size + 1
+    payload_without_digest = self.pack(link_protocol)[header_size:]
+    new_digest.update(payload_without_digest)
+
+    # Pack a copy of ourselves with our newly calculated digest, and encrypt
+    # the payload. Header remains plaintext.
+
+    cell = RelayCell(self.circ_id, self.command, self.data, new_digest, self.stream_id, self.recognized, self.unused)
+    header, payload = split(cell.pack(link_protocol), header_size)
+
+    return header + new_key.update(payload), new_key, new_digest
+
   @classmethod
   def _unpack(cls, content, circ_id, link_protocol):
     command, content = Size.CHAR.pop(content)
diff --git a/stem/client/datatype.py b/stem/client/datatype.py
index 5ca4e820..644d4993 100644
--- a/stem/client/datatype.py
+++ b/stem/client/datatype.py
@@ -116,6 +116,7 @@ import collections
 import hashlib
 import struct
 
+import stem.client.cell
 import stem.prereq
 import stem.util
 import stem.util.connection
@@ -246,9 +247,11 @@ class LinkProtocol(int):
     protocol = int.__new__(cls, version)
     protocol.version = version
     protocol.circ_id_size = Size.LONG if version > 3 else Size.SHORT
-    protocol.fixed_cell_length = 514 if version > 3 else 512
     protocol.first_circ_id = 0x80000000 if version > 3 else 0x01
 
+    cell_header_size = protocol.circ_id_size.size + 1  # circuit id (2 or 4 bytes) + command (1 byte)
+    protocol.fixed_cell_length = cell_header_size + stem.client.cell.FIXED_PAYLOAD_LEN
+
     return protocol
 
   def __hash__(self):
@@ -355,10 +358,15 @@ class Size(Field):
     raise NotImplementedError("Use our constant's unpack() and pop() instead")
 
   def pack(self, content):
-    if not stem.util._is_int(content):
-      raise ValueError('Size.pack encodes an integer, but was a %s' % type(content).__name__)
-
-    packed = struct.pack(self.format, content)
+    try:
+      packed = struct.pack(self.format, content)
+    except struct.error:
+      if not stem.util._is_int(content):
+        raise ValueError('Size.pack encodes an integer, but was a %s' % type(content).__name__)
+      elif content < 0:
+        raise ValueError('Packed values must be positive (attempted to pack %i as a %s)' % (content, self.name))
+      else:
+        raise  # some other struct exception
 
     if self.size != len(packed):
       raise ValueError('%s is the wrong size for a %s field' % (repr(packed), self.name))
diff --git a/stem/util/log.py b/stem/util/log.py
index 81fc88a6..0b2de90a 100644
--- a/stem/util/log.py
+++ b/stem/util/log.py
@@ -103,7 +103,7 @@ def get_logger():
   """
   Provides the stem logger.
 
-  :return: **logging.Logger** for stem
+  :returns: **logging.Logger** for stem
   """
 
   return LOGGER
diff --git a/stem/util/term.py b/stem/util/term.py
index 3c6fe73a..b4acd61f 100644
--- a/stem/util/term.py
+++ b/stem/util/term.py
@@ -80,7 +80,7 @@ def encoding(*attrs):
     :data:`~stem.util.terminal.BgColor`, or :data:`~stem.util.terminal.Attr` to
     provide an ecoding for
 
-  :return: **str** of the ANSI escape sequence, **None** no attributes are
+  :returns: **str** of the ANSI escape sequence, **None** no attributes are
     recognized
   """
 
diff --git a/test/unit/client/cell.py b/test/unit/client/cell.py
index ce492638..22843799 100644
--- a/test/unit/client/cell.py
+++ b/test/unit/client/cell.py
@@ -205,6 +205,7 @@ class TestCell(unittest.TestCase):
       self.assertEqual(digest, cell.digest)
       self.assertEqual(stream_id, cell.stream_id)
       self.assertEqual(unused, cell.unused)
+
       self.assertEqual(cell_bytes, cell.pack(link_protocol))
 
     digest = hashlib.sha1(b'hi')
diff --git a/test/unit/client/size.py b/test/unit/client/size.py
index eebe3619..dcd5fea3 100644
--- a/test/unit/client/size.py
+++ b/test/unit/client/size.py
@@ -25,6 +25,7 @@ class TestSize(unittest.TestCase):
     self.assertEqual(b'\x00\x00\x00\x00\x00\x00\x00\x12', Size.LONG_LONG.pack(18))
 
     self.assertRaisesWith(ValueError, 'Size.pack encodes an integer, but was a str', Size.CHAR.pack, 'hi')
+    self.assertRaisesWith(ValueError, 'Packed values must be positive (attempted to pack -1 as a CHAR)', Size.CHAR.pack, -1)
 
     bad_size = Size('BAD_SIZE', 1, '!H')
     self.assertRaisesRegexp(ValueError, re.escape("'\\x00\\x12' is the wrong size for a BAD_SIZE field"), bad_size.pack, 18)
diff --git a/test/unit/control/controller.py b/test/unit/control/controller.py
index 7ee2ab70..ba55208a 100644
--- a/test/unit/control/controller.py
+++ b/test/unit/control/controller.py
@@ -5,7 +5,6 @@ integ tests, but a few bits lend themselves to unit testing.
 
 import datetime
 import io
-import time
 import unittest
 
 import stem.descriptor.router_status_entry