[tor-commits] [stem/master] Add kdf_tor function

[atagar at torproject.org]

commit 73ba1f3506d44e8d65c439b404f7cd40eed5ceb9
Author: Damian Johnson <atagar at torproject.org>
Date:   Tue Jan 30 11:35:25 2018 -0800

    Add kdf_tor function
    
    Supporting KDF-TOR for the CREATE_FAST handshake. Test data is derived from
    what endosome's function provides.
---
 stem/client/__init__.py      | 24 ++++++++++++++++++++++++
 stem/client/cell.py          |  3 +--
 test/settings.cfg            |  1 +
 test/unit/client/__init__.py |  1 +
 test/unit/client/kdf_tor.py  | 19 +++++++++++++++++++
 5 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/stem/client/__init__.py b/stem/client/__init__.py
index fb0f899d..663bc048 100644
--- a/stem/client/__init__.py
+++ b/stem/client/__init__.py
@@ -12,6 +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
 
   Field - Packable and unpackable datatype.
     |- Size - Field of a static size.
@@ -106,6 +107,7 @@ a wrapper for :class:`~stem.socket.RelaySocket`, much the same way as
   ===================== ===========
 """
 
+import hashlib
 import io
 import struct
 
@@ -115,6 +117,8 @@ import stem.util.enum
 from stem.util import _hash_attr
 
 ZERO = '\x00'
+HASH_LEN = 20
+KEY_LEN = 16
 
 __all__ = [
   'cell',
@@ -457,6 +461,26 @@ class Certificate(Field):
     return _hash_attr(self, 'type_int', 'value')
 
 
+def kdf_tor(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)
+
+  :returns: **bytes** with the KDF-TOR of the key
+  """
+
+  derived_key_len = KEY_LEN * 2 + HASH_LEN * 3
+  derived_key, counter = '', 0
+
+  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]
+
+
 setattr(Size, 'CHAR', Size('CHAR', 1, '!B'))
 setattr(Size, 'SHORT', Size('SHORT', 2, '!H'))
 setattr(Size, 'LONG', Size('LONG', 4, '!L'))
diff --git a/stem/client/cell.py b/stem/client/cell.py
index 59fbc14c..9b0c3274 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -47,12 +47,11 @@ import sys
 import stem.client
 
 from stem import UNDEFINED
-from stem.client import ZERO, Address, Size, split
+from stem.client import HASH_LEN, ZERO, Address, Size, split
 from stem.util import _hash_attr, datetime_to_unix
 
 FIXED_PAYLOAD_LEN = 509
 AUTH_CHALLENGE_SIZE = 32
-HASH_LEN = 20
 
 
 class Cell(object):
diff --git a/test/settings.cfg b/test/settings.cfg
index 518ae2b2..7af57181 100644
--- a/test/settings.cfg
+++ b/test/settings.cfg
@@ -233,6 +233,7 @@ test.unit_tests
 |test.unit.client.size.TestSize
 |test.unit.client.address.TestAddress
 |test.unit.client.certificate.TestCertificate
+|test.unit.client.kdf_tor.TestKdfTor
 |test.unit.client.cell.TestCell
 |test.unit.connection.authentication.TestAuthenticate
 |test.unit.connection.connect.TestConnect
diff --git a/test/unit/client/__init__.py b/test/unit/client/__init__.py
index 7e745352..ba93b3e6 100644
--- a/test/unit/client/__init__.py
+++ b/test/unit/client/__init__.py
@@ -6,6 +6,7 @@ __all__ = [
   'address',
   'cell',
   'certificate',
+  'kdf_tor',
   'size',
 ]
 
diff --git a/test/unit/client/kdf_tor.py b/test/unit/client/kdf_tor.py
new file mode 100644
index 00000000..894d4fae
--- /dev/null
+++ b/test/unit/client/kdf_tor.py
@@ -0,0 +1,19 @@
+"""
+Unit tests for stem.client.kdf_tor.
+"""
+
+import unittest
+
+import stem.client
+
+KEY_1 = '\xec\xec.\xeb7R\xf2\n\xcb\xce\x97\xf4\x86\x82\x19#\x10\x0f\x08\xf0\xa2Z\xdeJ\x8f2\x8cc\xf6\xfa\x0e\t\x83f\xc5\xe2\xb3\x94\xa8\x13'
+KEY_2 = '\xe0v\xe4\xfaTB\x91\x1c\x81Gz\xa0\tI\xcb{\xc56\xcfV\xc2\xa0\x19\x9c\x98\x9a\x06\x0e\xc5\xfa\xb0z\x83\xa6\x10\xf6r"<b'
+
+DERIVED_1 = '\xca+\x81\x05\x14\x9d)o\xa6\x82\xe9B\xa8?\xf2\xaf\x85\x1b]6\xac\xcc\xbc\x91\xb1\xaf\xd7\xe0\xe9\x9dF#\xd8\xdbz\xe8\xe6\xca\x83,*\xe5scX\xbb+\xca \xcb\xa4\xbc\xad\x0f\x95\x0cO\xcc\xac\xf1\xc3\xbe\xc9\xe1\xf4\x90f\xdai\xf3\xf3\xf5\x14\xb5\xb9\x03U\xaf\x1e\x1b\xb1q||\x86A<_\xf7\xa0%\x86'
+DERIVED_2 = '\xbc0\xf99\x8e;Te\xbb+\xdb\xabR3l\xb9f?\x07KZC8\xe7\xa15\xd1IS\xd9\xd4\x1e\x96\xf6\xcd\x82\x91\x0b}r\x7f\xc5\xc0\xb1/[\x97dW\xba\x82g\xe7m^\x06[\xe6\xf8\xb4\x83f>c\x8b\x0f\x03\xcc\x98\x1f~t\x88\xe1\x83\xec\xbf*_\x8cF\x0e1\xa9\x17\xce\xa6\xa3\xd1+]\x1f'
+
+
+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))