[tor-commits] [stem/master] New RelayCell encrypt method

[atagar at torproject.org]

commit 8e83553e578aa0f8819eead434ca7ca26d57d5d9
Author: Damian Johnson <atagar at torproject.org>
Date:   Sun Aug 26 11:35:37 2018 -0700

    New RelayCell encrypt method
    
    "Premature optimization is the root of all evil."
    
    While I love the idea of moving our encrypt() and decrypt() methods into the
    RelayCell class I gotta admit I'm not a fan of the verbosity and large number
    of helpers. This branche's new RelayCell is far, far more complicated without
    adding any new capabilities to the Circuit class.
    
    As such basing this encrypt() method on our previous Circuit send() code.
    Hopefully this gets us the best of both worlds: Circuit is simpler because
    it no longer does encryption/decryption itself, and our RelayCell remains
    an elegantly simple class too.
    
    I suspect Dave might be factoring things this way in anticipation of future
    additions but I'd rather not complicate our code until it comes hand-in-hand
    with a practical benefit in doing so.
---
 stem/client/__init__.py |  2 +-
 stem/client/cell.py     | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 36 insertions(+), 1 deletion(-)

diff --git a/stem/client/__init__.py b/stem/client/__init__.py
index 19528d67..486ff0a6 100644
--- a/stem/client/__init__.py
+++ b/stem/client/__init__.py
@@ -237,7 +237,7 @@ class Circuit(object):
       # successfully sent.
 
       cell = stem.client.cell.RelayCell(self.id, command, data, stream_id = stream_id)
-      payload, forward_digest, forward_key = cell.encrypt(self.relay.link_protocol, self.forward_digest, self.forward_key)
+      payload, forward_digest, forward_key = cell.encrypt(self.relay.link_protocol, self.forward_key, self.forward_digest)
       self.relay._orport.send(payload)
 
       self.forward_digest = forward_digest
diff --git a/stem/client/cell.py b/stem/client/cell.py
index b40a30e1..29aa18ab 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -576,6 +576,41 @@ class RelayCell(CircuitCell):
 
     return RelayCell._pack(link_protocol, bytes(payload), self.unused, self.circ_id)
 
+  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.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)