[tor-commits] [stem/master] Replace kdf_tor function with a KDF class

[atagar at torproject.org]

commit e2e7b5d58fdc14014df8c476c59457b1b5483c2e
Author: Damian Johnson <atagar at torproject.org>
Date:   Tue Jan 30 12:21:04 2018 -0800

    Replace kdf_tor function with a KDF class
    
    On reflection callers actually don't care about the KDF-TOR value. Rather, they
    care about its derived attributes.
---
 stem/client/__init__.py     | 47 ++++++++++++++++++++++++++++++++-------------
 test/unit/client/kdf_tor.py | 21 +++++++++++++++++---
 2 files changed, 52 insertions(+), 16 deletions(-)

diff --git a/stem/client/__init__.py b/stem/client/__init__.py
index 663bc048..4c6bdb12 100644
--- a/stem/client/__init__.py
+++ b/stem/client/__init__.py
@@ -12,7 +12,7 @@ a wrapper for :class:`~stem.socket.RelaySocket`, much the same way as
 ::
 
   split - splits bytes into substrings
-  kdf_tor - calculates the derived key using the KDF-TOR protocol
+  KDF - KDF-TOR key derivative for TAP, CREATE_FAST handshakes, and hidden serivces
 
   Field - Packable and unpackable datatype.
     |- Size - Field of a static size.
@@ -461,24 +461,45 @@ class Certificate(Field):
     return _hash_attr(self, 'type_int', 'value')
 
 
-def kdf_tor(key):
+class KDF(object):
+  """
+  Tor's derived key for TAP, CREATE_FAST handshakes, and hidden service
+  protocols as defined tor-spec section 5.2.1.
+
+  :var bytes key_hash: expected derived key that proves knowledge of our shared
+    computed key
+  :var bytes forward_digest: forward digest hash seed
+  :var bytes backward_digest: backward digest hash seed
+  :var bytes forward_key: forward encryption key
+  :var bytes backward_key: backward encryption key
   """
-  Tor's key derivation function used by TAP, CREATE_FAST handshakes, and hidden
-  service protocols as defined in section 5.2.1 of the tor spec.
 
-  :param bytes key: shared key with endpoint (K0 in the spec)
+  def __init__(self, key_material):
+    value = KDF._value(key_material)
 
-  :returns: **bytes** with the KDF-TOR of the key
-  """
+    self.key_hash, value = split(value, HASH_LEN)
+    self.forward_digest, value = split(value, HASH_LEN)
+    self.backward_digest, value = split(value, HASH_LEN)
+    self.forward_key, value = split(value, KEY_LEN)
+    self.backward_key, value = split(value, KEY_LEN)
+
+  @staticmethod
+  def _value(key):
+    """
+    Computes the KDF-TOR value...
+
+      K = H(K0 | [00]) | H(K0 | [01]) | H(K0 | [02]) | ...
+    """
 
-  derived_key_len = KEY_LEN * 2 + HASH_LEN * 3
-  derived_key, counter = '', 0
+    derived_key = ''
+    derived_key_len = KEY_LEN * 2 + HASH_LEN * 3
+    counter = 0
 
-  while len(derived_key) < derived_key_len:
-    derived_key += hashlib.sha1(key + Size.CHAR.pack(counter)).digest()
-    counter += 1
+    while len(derived_key) < derived_key_len:
+      derived_key += hashlib.sha1(key + Size.CHAR.pack(counter)).digest()
+      counter += 1
 
-  return derived_key[:derived_key_len]
+    return derived_key[:derived_key_len]
 
 
 setattr(Size, 'CHAR', Size('CHAR', 1, '!B'))
diff --git a/test/unit/client/kdf_tor.py b/test/unit/client/kdf_tor.py
index 894d4fae..f356d2cd 100644
--- a/test/unit/client/kdf_tor.py
+++ b/test/unit/client/kdf_tor.py
@@ -14,6 +14,21 @@ DERIVED_2 = '\xbc0\xf99\x8e;Te\xbb+\xdb\xabR3l\xb9f?\x07KZC8\xe7\xa15\xd1IS\xd9\
 
 
 class TestKdfTor(unittest.TestCase):
-  def test_kdf_tor(self):
-    self.assertEqual(DERIVED_1, stem.client.kdf_tor(KEY_1))
-    self.assertEqual(DERIVED_2, stem.client.kdf_tor(KEY_2))
+  def test_kdf_value(self):
+    self.assertEqual(DERIVED_1, stem.client.KDF._value(KEY_1))
+    self.assertEqual(DERIVED_2, stem.client.KDF._value(KEY_2))
+
+  def test_kdf_attributes(self):
+    k1 = stem.client.KDF(KEY_1)
+    self.assertEqual('\xca+\x81\x05\x14\x9d)o\xa6\x82\xe9B\xa8?\xf2\xaf\x85\x1b]6', k1.key_hash)
+    self.assertEqual('\xac\xcc\xbc\x91\xb1\xaf\xd7\xe0\xe9\x9dF#\xd8\xdbz\xe8\xe6\xca\x83,', k1.forward_digest)
+    self.assertEqual('*\xe5scX\xbb+\xca \xcb\xa4\xbc\xad\x0f\x95\x0cO\xcc\xac\xf1', k1.backward_digest)
+    self.assertEqual('\xc3\xbe\xc9\xe1\xf4\x90f\xdai\xf3\xf3\xf5\x14\xb5\xb9\x03', k1.forward_key)
+    self.assertEqual('U\xaf\x1e\x1b\xb1q||\x86A<_\xf7\xa0%\x86', k1.backward_key)
+
+    k2 = stem.client.KDF(KEY_1)
+    self.assertEqual('\xca+\x81\x05\x14\x9d)o\xa6\x82\xe9B\xa8?\xf2\xaf\x85\x1b]6', k2.key_hash)
+    self.assertEqual('\xac\xcc\xbc\x91\xb1\xaf\xd7\xe0\xe9\x9dF#\xd8\xdbz\xe8\xe6\xca\x83,', k2.forward_digest)
+    self.assertEqual('*\xe5scX\xbb+\xca \xcb\xa4\xbc\xad\x0f\x95\x0cO\xcc\xac\xf1', k2.backward_digest)
+    self.assertEqual('\xc3\xbe\xc9\xe1\xf4\x90f\xdai\xf3\xf3\xf5\x14\xb5\xb9\x03', k2.forward_key)
+    self.assertEqual('U\xaf\x1e\x1b\xb1q||\x86A<_\xf7\xa0%\x86', k2.backward_key)