[tor-commits] [stem/master] Add decrypt instance method to BaseRelayCell

[atagar at torproject.org]

commit 5baa2e37728ab50a5132d81225070e097e6bd058
Author: Dave Rolek <dmr-x at riseup.net>
Date:   Sat Aug 18 03:58:13 2018 +0000

    Add decrypt instance method to BaseRelayCell
    
    Akin to encrypt(), this takes care of decryption and all the ancillary
    functionality related to it.
    
    It composes the functionality of the building blocks added in previous
    changes, namely:
      * check_recognized_field()
      * check_digest()
      * interpret_cell() - but not automatically done by default
    
    A consumer of this, especially in a multi-hop circuit, will only need to
    manage the digest/decryptor state (in the case of a cell that cannot be
    decrypted+recognized... fallback) and pay attention to the
    'fully_decrypted' result.
    
    This should make a consumer's implementation pretty readable, something
    like:
    
      # relay_1 := nearest relay in circuit
      decryption_states = [
        (relay_1, digest_1, decryptor_1),
        (relay_2, digest_2, decryptor_2),
        (relay_3, digest_3, decryptor_3),
      ]
    
      new_decryption_states = copy.deepcopy(decryption_states)
    
      from_relay = None
      for i in range(len(new_decryption_states)):
        relay, digest, decryptor = new_decryption_states[i]
    
        cell, fully_decrypted, digest, decryptor = cell.decrypt(digest, decryptor)
        new_decryption_states[i] = (relay, digest, decryptor)
    
        if fully_decrypted:
          from_relay = relay
          break
    
      if from_relay:
        decryption_states = new_decryption_states
      else:
        # handle this, since the cell couldn't be decrypted
        # probably raise Something()
        pass
    
      cell_and_hop = (cell, from_relay)  # naming things is hard
    
      # do something with the decrypted cell
      # ^ probably feeding it into the queue/handler for the (circuit_id, relay, stream_id)
---
 stem/client/cell.py | 65 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/stem/client/cell.py b/stem/client/cell.py
index a96b2f7b..48dbcf36 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -443,6 +443,71 @@ class BaseRelayCell(CircuitCell):
 
     return new_cell
 
+  def decrypt(self, digest, decryptor, interpret = False):
+    """
+    Decrypts a cell and checks whether it is fully decrypted,
+    returning a new (Cell, fully_decrypted, digest, decryptor) tuple.
+    Optionally also interprets the cell (not generally recommended).
+
+    The method name is technically a misnomer, as it also checks whether the
+    cell has been fully decrypted (after decrypting), updating the digest if so.
+    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
+          do all of that, for pedantry.
+
+    :param HASH digest: running digest held with the relay
+    :param cryptography.hazmat.primitives.ciphers.CipherContext decryptor:
+      running stream cipher decryptor held with the relay
+
+    :param bool interpret: (optional, defaults to **False**) Use **True** with
+      caution. The spec indicates that a fully decrypted cell should be
+      accounted for in digest and decryptor, independent of cell validity. Using
+      **True**, while convenient, may cause an exception for a NYI relay
+      command, a malformed cell, or some other reason. This option should only
+      be used when the consumer will consider the circuit to have a fatal error
+      in such cases, and catches/handles the exception accordingly (e.g. sending
+      a DestroyCell).
+
+    :returns: (:class:`~stem.client.cell.Cell`, bool, HASH, CipherContext) tuple
+      of object copies updated as follows:
+        * Cell: either :class:`~stem.client.cell.RawRelayCell` with a decrypted
+          payload or :class:`~stem.client.cell.RelayCell` class or subclass, if
+          **interpret** is **True** and the cell was fully decrypted
+        * fully_decrypted: **bool** indicating whether the cell is fully
+          decrypted
+        * digest: updated via digest.update(payload), if the cell was fully
+          decrypted; otherwise a copy of the original
+        * decryptor: updated via decryptor.update(payload)
+    """
+
+    new_decryptor = copy.copy(decryptor)
+
+    # actually decrypt
+    decrypted_payload = new_decryptor.update(self.payload)
+    new_cell = self.__class__(self.circ_id, decrypted_payload)
+
+    # do post-decryption checks to ascertain whether cell is fully decrypted
+    if new_cell.check_recognized_field():
+      digest_matches, new_digest = new_cell.check_digest(digest)
+      fully_decrypted = digest_matches
+    else:
+      new_digest = None
+      fully_decrypted = False
+
+    # only return the new_digest if the digest check meant that the cell has been fully decrypted
+    #
+    # furthermore, even if the digest was not updated, return a copy
+    # this allows a consumer to always assume the returned digest is a different object
+    digest_to_return = new_digest if fully_decrypted else digest.copy()
+
+    if interpret and fully_decrypted:
+      # this might raise an exception; oh well, we did warn about that
+      new_cell = new_cell.interpret_cell()
+
+    return new_cell, fully_decrypted, digest_to_return, new_decryptor
+
   def __hash__(self):
     return stem.util._hash_attr(self, 'circ_id', 'payload', cache = True)