[tor-commits] [stem/master] Implement AUTH_CHALLENGE cells

[atagar at torproject.org]

commit 06767ca7c64a3d399e75cb92610cf80197fe7039
Author: Damian Johnson <atagar at torproject.org>
Date:   Sun Jan 14 17:31:37 2018 -0800

    Implement AUTH_CHALLENGE cells
---
 stem/client/cell.py      | 61 ++++++++++++++++++++++++++++++++++++++++++++++--
 test/unit/client/cell.py | 19 ++++++++++++++-
 2 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/stem/client/cell.py b/stem/client/cell.py
index 872f0998..df34508d 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -46,6 +46,7 @@ from stem import UNDEFINED
 from stem.client import ZERO, Certificate, Size
 
 FIXED_PAYLOAD_LEN = 509
+AUTH_CHALLENGE_SIZE = 32
 
 
 class Cell(object):
@@ -444,18 +445,74 @@ class CertsCell(Cell):
       if cert_size > len(content):
         raise ValueError('CERTS cell should have a certificate with %i bytes, but only had %i remaining' % (cert_size, len(content)))
 
-      cert_bytes = content[:cert_size]
-      content = content[cert_size:]
+      cert_bytes, content = content[:cert_size], content[cert_size:]
       certs.append(Certificate(cert_type, cert_bytes))
 
     return CertsCell(certs)
 
 
 class AuthChallengeCell(Cell):
+  """
+  First step of the authentication handshake.
+
+  :var bytes challenge: random bytes for us to sign to authenticate
+  :var list methods: authentication methods supported by the relay we're
+    communicating with
+  """
+
   NAME = 'AUTH_CHALLENGE'
   VALUE = 130
   IS_FIXED_SIZE = False
 
+  def __init__(self, challenge, methods):
+    self.challenge = challenge
+    self.methods = methods
+
+  @classmethod
+  def pack(cls, link_version, methods, challenge = None):
+    """
+    Provides an authentication challenge.
+
+    :param int link_version: link protocol version
+    :param list methods: authentication methods we support
+    :param bytes challenge: randomized string for the receiver to sign
+
+    :returns: **bytes** with a payload for this challenge
+    """
+
+    if challenge is None:
+      challenge = os.urandom(AUTH_CHALLENGE_SIZE)
+    elif len(challenge) != AUTH_CHALLENGE_SIZE:
+      raise ValueError('AUTH_CHALLENGE must be %i bytes, but was %i' % (AUTH_CHALLENGE_SIZE, len(challenge)))
+
+    payload = io.BytesIO()
+    payload.write(challenge)
+    payload.write(Size.SHORT.pack(len(methods)))
+
+    for method in methods:
+      payload.write(Size.SHORT.pack(method))
+
+    return cls._pack(link_version, payload.getvalue())
+
+  @classmethod
+  def _unpack(cls, content, circ_id, link_version):
+    if len(content) < AUTH_CHALLENGE_SIZE + 2:
+      raise ValueError('AUTH_CHALLENGE payload should be at least 34 bytes, but was %i' % len(content))
+
+    challenge, content = content[:AUTH_CHALLENGE_SIZE], content[AUTH_CHALLENGE_SIZE:]
+    method_count, content = Size.SHORT.pop(content)
+
+    if len(content) < method_count * 2:
+      raise ValueError('AUTH_CHALLENGE should have %i methods, but only had %i bytes for it' % (method_count, len(content)))
+
+    methods = []
+
+    for i in range(method_count):
+      method, content = Size.SHORT.pop(content)
+      methods.append(method)
+
+    return AuthChallengeCell(challenge, methods)
+
 
 class AuthenticateCell(Cell):
   NAME = 'AUTHENTICATE'
diff --git a/test/unit/client/cell.py b/test/unit/client/cell.py
index 0d2b677d..f95a664e 100644
--- a/test/unit/client/cell.py
+++ b/test/unit/client/cell.py
@@ -15,9 +15,11 @@ from stem.client.cell import (
   VersionsCell,
   VPaddingCell,
   CertsCell,
+  AuthChallengeCell,
 )
 
 RANDOM_PAYLOAD = os.urandom(FIXED_PAYLOAD_LEN)
+CHALLENGE = '\x89Y\t\x99\xb2\x1e\xd9*V\xb6\x1bn\n\x05\xd8/\xe3QH\x85\x13Z\x17\xfc\x1c\x00{\xa9\xae\x83^K'
 
 PADDING_CELLS = {
   '\x00\x00\x00' + RANDOM_PAYLOAD: RANDOM_PAYLOAD,
@@ -42,6 +44,10 @@ CERTS_CELLS = {
   '\x00\x00\x81\x00\x05\x01\x01\x00\x01\x08': [Certificate(type = 1, value = '\x08')],
 }
 
+AUTH_CHALLENGE_CELLS = {
+  '\x00\x00\x82\x00&%s\x00\x02\x00\x01\x00\x03' % CHALLENGE: (CHALLENGE, [1, 3]),
+}
+
 
 class TestCell(unittest.TestCase):
   def test_by_name(self):
@@ -69,7 +75,7 @@ class TestCell(unittest.TestCase):
 
   def test_unpack_for_new_link(self):
     # TODO: we need to support more cell types before we can test this
-    self.assertRaisesRegexp(NotImplementedError, 'Unpacking not yet implemented for AUTH_CHALLENGE cells', Cell.unpack, test_data('new_link_cells'), 2)
+    self.assertRaisesRegexp(NotImplementedError, 'Unpacking not yet implemented for NETINFO cells', Cell.unpack, test_data('new_link_cells'), 2)
 
   def test_padding_packing(self):
     for cell_bytes, payload in PADDING_CELLS.items():
@@ -101,3 +107,14 @@ class TestCell(unittest.TestCase):
 
     self.assertRaisesRegexp(ValueError, 'CERTS cell should have a certificate with 3 bytes, but only had 1 remaining', Cell.unpack, '\x00\x00\x81\x00\x05\x01\x01\x00\x03\x08', 2)
     self.assertRaisesRegexp(ValueError, 'CERTS cell indicates it should have 2 certificates, but only contained 1', Cell.unpack, '\x00\x00\x81\x00\x05\x02\x01\x00\x01\x08', 2)
+
+  def test_auth_challenge_packing(self):
+    for cell_bytes, (challenge, methods) in AUTH_CHALLENGE_CELLS.items():
+      self.assertEqual(cell_bytes, AuthChallengeCell.pack(2, methods, challenge))
+
+      cell = Cell.unpack(cell_bytes, 2)[0]
+      self.assertEqual(challenge, cell.challenge)
+      self.assertEqual(methods, cell.methods)
+
+    self.assertRaisesRegexp(ValueError, 'AUTH_CHALLENGE cell should have a payload of 38 bytes, but only had 16', Cell.unpack, '\x00\x00\x82\x00&%s\x00\x02\x00\x01\x00\x03' % CHALLENGE[:10], 2)
+    self.assertRaisesRegexp(ValueError, 'AUTH_CHALLENGE should have 3 methods, but only had 4 bytes for it', Cell.unpack, '\x00\x00\x82\x00&%s\x00\x03\x00\x01\x00\x03' % CHALLENGE, 2)