[tor-commits] [stem/master] Refactor RELAY cell encryption into new encrypt instance method

[atagar at torproject.org]

commit 116eb4b3efd01a91093a81d9411582cd43c98584
Author: Dave Rolek <dmr-x at riseup.net>
Date:   Fri Aug 17 19:30:56 2018 +0000

    Refactor RELAY cell encryption into new encrypt instance method
    
    This further tunes the abstraction layers such that consumers don't need
    to care at all what part of the Cell is encrypted. They are still
    responsible for managing the digest/encryption key states, but don't
    need to know the finer details of the encryption.
---
 stem/client/__init__.py |  6 +-----
 stem/client/cell.py     | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/stem/client/__init__.py b/stem/client/__init__.py
index b64956ff..a228bebc 100644
--- a/stem/client/__init__.py
+++ b/stem/client/__init__.py
@@ -239,11 +239,7 @@ class Circuit(object):
 
       try:
         cell = stem.client.cell.RelayCell(self.id, command, data, stream_id = stream_id)
-        cell_with_digest, self.forward_digest = cell.apply_digest(self.forward_digest)
-        payload_with_digest = cell_with_digest.pack_payload()
-
-        encrypted_payload = self.forward_key.update(payload_with_digest)
-        encrypted_cell = stem.client.cell.RawRelayCell(self.id, encrypted_payload)
+        encrypted_cell, self.forward_digest, self.forward_key = cell.encrypt(self.forward_digest, self.forward_key)
 
         reply_cells = []
         self.relay._orport.send(encrypted_cell.pack(self.relay.link_protocol))
diff --git a/stem/client/cell.py b/stem/client/cell.py
index d1265c10..43dc076c 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -39,6 +39,7 @@ Messages communicated over a Tor relay's ORPort.
     +- pop - decodes cell with remainder
 """
 
+import copy
 import datetime
 import inspect
 import os
@@ -499,6 +500,40 @@ class RelayCell(CircuitCell):
 
     return new_cell, new_digest
 
+  def encrypt(self, digest, encryptor, **kwargs):
+    """
+    Preps a cell payload, including calculating digest, and encrypts it,
+    returning a new (RawRelayCell, digest, encryptor) tuple.
+
+    The method name is technically a misnomer, as it also preps cell payload
+    and applies the digest, prior to encrypting. However, these operations
+    are defined per the spec as required for RELAY cells, and ...
+      (1) it is a natural mental extension to include them here;
+      (2) it would be a bit pointless to require method consumers to manually
+          call both, for pedantry.
+
+    :param HASH digest: running digest held with the relay
+    :param cryptography.hazmat.primitives.ciphers.CipherContext encryptor:
+      running stream cipher encryptor held with the relay
+
+    :param bool prep_cell: (optional, defaults to **True**) refer to
+      :func:`~stem.client.cell.RelayCell.apply_digest`
+
+    :returns: (:class:`~stem.client.cell.RawRelayCell`, HASH, CipherContext)
+      tuple of object copies updated as follows:
+        * RawRelayCell: updated as specified in
+          :func:`~stem.client.cell.RelayCell.apply_digest`, then encrypted
+        * digest: updated via digest.update(payload)
+        * encryptor: updated via encryptor.update(payload_with_digest)
+    """
+
+    unencrypted_cell, new_digest = self.apply_digest(digest, **kwargs)
+    new_encryptor = copy.copy(encryptor)
+    encrypted_payload = new_encryptor.update(unencrypted_cell.pack_payload())
+    encrypted_cell = RawRelayCell(unencrypted_cell.circ_id, encrypted_payload)
+
+    return encrypted_cell, new_digest, new_encryptor
+
   def pack_payload(self, **kwargs):
     """
     Convenience method for running