[tor-commits] [stem/master] Initial Circuit.send() method

[atagar at torproject.org]

commit b6ec30d450e929976acf3fce9378f0c3a777d938
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Feb 5 11:24:34 2018 -0800

    Initial Circuit.send() method
    
    Quite a few things I dislike about this, but it works! Next up: testing and
    refinement...
---
 stem/client/__init__.py | 37 ++++++++++++++++++++++++++++++++++++-
 stem/client/cell.py     | 17 +++++++++++++++--
 2 files changed, 51 insertions(+), 3 deletions(-)

diff --git a/stem/client/__init__.py b/stem/client/__init__.py
index fdb436b6..4f9ed17c 100644
--- a/stem/client/__init__.py
+++ b/stem/client/__init__.py
@@ -19,6 +19,7 @@ a wrapper for :class:`~stem.socket.RelaySocket`, much the same way as
     +- close - shuts down our connection
 """
 
+import copy
 import hashlib
 
 import stem
@@ -26,7 +27,7 @@ import stem.client.cell
 import stem.socket
 import stem.util.connection
 
-from stem.client.datatype import ZERO, AddrType, Address, KDF
+from stem.client.datatype import ZERO, AddrType, Address, KDF, split
 
 __all__ = [
   'cell',
@@ -207,3 +208,37 @@ class Circuit(object):
     self.backward_digest = hashlib.sha1(kdf.backward_digest)
     self.forward_key = Cipher(algorithms.AES(kdf.forward_key), ctr, default_backend()).encryptor()
     self.backward_key = Cipher(algorithms.AES(kdf.backward_key), ctr, default_backend()).decryptor()
+
+  def send(self, command, data, stream_id = 0):
+    """
+    Sends a message over the circuit.
+
+    :param stem.client.RelayCommand command: command to be issued
+    :param bytes data: message payload
+    :param int stream_id: specific stream this concerns
+    """
+
+    # TODO: move RelayCommand to this base module?
+    # TODO: add lock
+
+    orig_digest = self.forward_digest.copy()
+    orig_key = copy.copy(self.forward_key)
+
+    try:
+      cell = stem.client.cell.RelayCell(self.id, command, data, 0, stream_id)
+      payload_without_digest = cell.pack(self.relay.link_protocol)[3:]
+      self.forward_digest.update(payload_without_digest)
+
+      cell = stem.client.cell.RelayCell(self.id, command, data, self.forward_digest, stream_id)
+      header, payload = split(cell.pack(self.relay.link_protocol), 3)
+      encrypted_payload = header + self.forward_key.update(payload)
+
+      self.relay._orport.send(encrypted_payload)
+    except:
+      self.forward_digest = orig_digest
+      self.forward_key = orig_key
+      raise
+
+  def close(self):
+    self.relay._orport.send(stem.client.cell.DestroyCell(self.id).pack(self.relay.link_protocol))
+    del self.relay._circuits[self.id]
diff --git a/stem/client/cell.py b/stem/client/cell.py
index 17f27202..c4a19940 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -307,7 +307,7 @@ class RelayCell(CircuitCell):
   VALUE = 3
   IS_FIXED_SIZE = True
 
-  def __init__(self, circ_id, command, data, digest = 0, stream_id = 0):
+  def __init__(self, circ_id, command, data, digest = 0, stream_id = 0, raw_content = None):
     if 'hashlib.HASH' in str(type(digest)):
       # Unfortunately hashlib generates from a dynamic private class so
       # isinstance() isn't such a great option.
@@ -325,6 +325,7 @@ class RelayCell(CircuitCell):
     self.data = data
     self.digest = digest
     self.stream_id = stream_id
+    self._raw_content = raw_content
 
     if not stream_id and self.command in STREAM_ID_REQUIRED:
       raise ValueError('%s relay cells require a stream id' % self.command)
@@ -342,8 +343,20 @@ class RelayCell(CircuitCell):
 
     return RelayCell._pack(link_protocol, payload.getvalue(), self.circ_id)
 
+  def decrypt(self, circ):
+    # TODO: clearly funky, just a spot to start...
+
+    if not self._raw_content:
+      raise ValueError('Only received cells can be decrypted')
+
+    decrypted = circ.backward_key.update(self._raw_content)
+    return RelayCell._unpack(decrypted, self.circ_id, 3)
+
+
   @classmethod
   def _unpack(cls, content, circ_id, link_protocol):
+    orig_content = content
+
     command, content = Size.CHAR.pop(content)
     _, content = Size.SHORT.pop(content)  # 'recognized' field
     stream_id, content = Size.SHORT.pop(content)
@@ -351,7 +364,7 @@ class RelayCell(CircuitCell):
     data_len, content = Size.SHORT.pop(content)
     data, content = split(content, data_len)
 
-    return RelayCell(circ_id, command, data, digest, stream_id)
+    return RelayCell(circ_id, command, data, digest, stream_id, orig_content)
 
   def __hash__(self):
     return _hash_attr(self, 'command_int', 'stream_id', 'digest', 'data')