[tor-commits] [stem/master] Supporting controllers in the connection module

atagar at torproject.org atagar at torproject.org
Sun May 6 01:13:58 UTC 2012


commit af691f5b54d1f571ee51d8c67010ea19e4c672fd
Author: Damian Johnson <atagar at torproject.org>
Date:   Sat May 5 17:48:29 2012 -0700

    Supporting controllers in the connection module
    
    Users will be using a BaseController subclass unless they really need to work
    at a low level, in which case they'll be using a ControlSocket. Making the
    connection module (which does authentication) support both.
---
 stem/connection.py                      |   86 +++++++++++++++++++------------
 test/integ/connection/authentication.py |   26 +++++++---
 test/integ/connection/connect.py        |    4 +-
 test/runner.py                          |   38 +++++++++++---
 4 files changed, 103 insertions(+), 51 deletions(-)

diff --git a/stem/connection.py b/stem/connection.py
index d950049..825edf5 100644
--- a/stem/connection.py
+++ b/stem/connection.py
@@ -277,7 +277,7 @@ def _connect(control_socket, password, chroot_path, controller):
     print "Unable to authenticate: %s" % exc
     return None
 
-def authenticate(control_socket, password = None, chroot_path = None, protocolinfo_response = None):
+def authenticate(controller, password = None, chroot_path = None, protocolinfo_response = None):
   """
   Authenticates to a control socket using the information provided by a
   PROTOCOLINFO response. In practice this will often be all we need to
@@ -288,7 +288,8 @@ def authenticate(control_socket, password = None, chroot_path = None, protocolin
   about, then have a AuthenticationFailure catch-all at the end.
   
   Arguments:
-    control_socket (stem.socket.ControlSocket) - socket to be authenticated
+    controller (stem.socket.ControlSocket or stem.control.BaseController) -
+      tor controller connection to be authenticated
     password (str) - passphrase to present to the socket if it uses password
         authentication (skips password auth if None)
     chroot_path (str) - path prefix if in a chroot environment
@@ -302,8 +303,8 @@ def authenticate(control_socket, password = None, chroot_path = None, protocolin
     follows...
     
     stem.connection.IncorrectSocketType
-      The control_socket does not speak the tor control protocol. Most often
-      this happened because the user confused the SocksPort or ORPort with the
+      The controller does not speak the tor control protocol. Most often this
+      happened because the user confused the SocksPort or ORPort with the
       ControlPort.
     
     stem.connection.UnrecognizedAuthMethods
@@ -355,7 +356,7 @@ def authenticate(control_socket, password = None, chroot_path = None, protocolin
   
   if not protocolinfo_response:
     try:
-      protocolinfo_response = get_protocolinfo(control_socket)
+      protocolinfo_response = get_protocolinfo(controller)
     except stem.socket.ProtocolError:
       raise IncorrectSocketType("unable to use the control socket")
     except stem.socket.SocketError, exc:
@@ -397,16 +398,16 @@ def authenticate(control_socket, password = None, chroot_path = None, protocolin
     
     try:
       if auth_type == AuthMethod.NONE:
-        authenticate_none(control_socket, False)
+        authenticate_none(controller, False)
       elif auth_type == AuthMethod.PASSWORD:
-        authenticate_password(control_socket, password, False)
+        authenticate_password(controller, password, False)
       elif auth_type == AuthMethod.COOKIE:
         cookie_path = protocolinfo_response.cookie_path
         
         if chroot_path:
           cookie_path = os.path.join(chroot_path, cookie_path.lstrip(os.path.sep))
         
-        authenticate_cookie(control_socket, cookie_path, False)
+        authenticate_cookie(controller, cookie_path, False)
       
       return # success!
     except OpenAuthRejected, exc:
@@ -439,7 +440,7 @@ def authenticate(control_socket, password = None, chroot_path = None, protocolin
   
   raise AssertionError("BUG: Authentication failed without providing a recognized exception: %s" % str(auth_exceptions))
 
-def authenticate_none(control_socket, suppress_ctl_errors = True):
+def authenticate_none(controller, suppress_ctl_errors = True):
   """
   Authenticates to an open control socket. All control connections need to
   authenticate before they can be used, even if tor hasn't been configured to
@@ -452,7 +453,8 @@ def authenticate_none(control_socket, suppress_ctl_errors = True):
   For general usage use the authenticate() function instead.
   
   Arguments:
-    control_socket (stem.socket.ControlSocket) - socket to be authenticated
+    controller (stem.socket.ControlSocket or stem.control.BaseController) -
+      tor controller connection
     suppress_ctl_errors (bool) - reports raised stem.socket.ControllerError as
       authentication rejection if True, otherwise they're re-raised
   
@@ -462,23 +464,22 @@ def authenticate_none(control_socket, suppress_ctl_errors = True):
   """
   
   try:
-    control_socket.send("AUTHENTICATE")
-    auth_response = control_socket.recv()
+    auth_response = _msg(controller, "AUTHENTICATE")
     
     # if we got anything but an OK response then error
     if str(auth_response) != "OK":
-      try: control_socket.connect()
+      try: controller.connect()
       except: pass
       
       raise OpenAuthRejected(str(auth_response), auth_response)
   except stem.socket.ControllerError, exc:
-    try: control_socket.connect()
+    try: controller.connect()
     except: pass
     
     if not suppress_ctl_errors: raise exc
     else: raise OpenAuthRejected("Socket failed (%s)" % exc)
 
-def authenticate_password(control_socket, password, suppress_ctl_errors = True):
+def authenticate_password(controller, password, suppress_ctl_errors = True):
   """
   Authenticates to a control socket that uses a password (via the
   HashedControlPassword torrc option). Quotes in the password are escaped.
@@ -495,7 +496,8 @@ def authenticate_password(control_socket, password, suppress_ctl_errors = True):
   future versions.
   
   Arguments:
-    control_socket (stem.socket.ControlSocket) - socket to be authenticated
+    controller (stem.socket.ControlSocket or stem.control.BaseController) -
+      tor controller connection
     password (str) - passphrase to present to the socket
     suppress_ctl_errors (bool) - reports raised stem.socket.ControllerError as
       authentication rejection if True, otherwise they're re-raised
@@ -514,12 +516,11 @@ def authenticate_password(control_socket, password, suppress_ctl_errors = True):
   password = password.replace('"', '\\"')
   
   try:
-    control_socket.send("AUTHENTICATE \"%s\"" % password)
-    auth_response = control_socket.recv()
+    auth_response = _msg(controller, "AUTHENTICATE \"%s\"" % password)
     
     # if we got anything but an OK response then error
     if str(auth_response) != "OK":
-      try: control_socket.connect()
+      try: controller.connect()
       except: pass
       
       # all we have to go on is the error message from tor...
@@ -531,13 +532,13 @@ def authenticate_password(control_socket, password, suppress_ctl_errors = True):
       else:
         raise PasswordAuthRejected(str(auth_response), auth_response)
   except stem.socket.ControllerError, exc:
-    try: control_socket.connect()
+    try: controller.connect()
     except: pass
     
     if not suppress_ctl_errors: raise exc
     else: raise PasswordAuthRejected("Socket failed (%s)" % exc)
 
-def authenticate_cookie(control_socket, cookie_path, suppress_ctl_errors = True):
+def authenticate_cookie(controller, cookie_path, suppress_ctl_errors = True):
   """
   Authenticates to a control socket that uses the contents of an authentication
   cookie (generated via the CookieAuthentication torrc option). This does basic
@@ -559,7 +560,8 @@ def authenticate_cookie(control_socket, cookie_path, suppress_ctl_errors = True)
   future versions.
   
   Arguments:
-    control_socket (stem.socket.ControlSocket) - socket to be authenticated
+    controller (stem.socket.ControlSocket or stem.control.BaseController) -
+      tor controller connection
     cookie_path (str) - path of the authentication cookie to send to tor
     suppress_ctl_errors (bool) - reports raised stem.socket.ControllerError as
       authentication rejection if True, otherwise they're re-raised
@@ -599,12 +601,12 @@ def authenticate_cookie(control_socket, cookie_path, suppress_ctl_errors = True)
     raise UnreadableCookieFile("Authentication failed: unable to read '%s' (%s)" % (cookie_path, exc), cookie_path) 
   
   try:
-    control_socket.send("AUTHENTICATE %s" % binascii.b2a_hex(auth_cookie_contents))
-    auth_response = control_socket.recv()
+    msg = "AUTHENTICATE %s" % binascii.b2a_hex(auth_cookie_contents)
+    auth_response = _msg(controller, msg)
     
     # if we got anything but an OK response then error
     if str(auth_response) != "OK":
-      try: control_socket.connect()
+      try: controller.connect()
       except: pass
       
       # all we have to go on is the error message from tor...
@@ -617,20 +619,21 @@ def authenticate_cookie(control_socket, cookie_path, suppress_ctl_errors = True)
       else:
         raise CookieAuthRejected(str(auth_response), cookie_path, auth_response)
   except stem.socket.ControllerError, exc:
-    try: control_socket.connect()
+    try: controller.connect()
     except: pass
     
     if not suppress_ctl_errors: raise exc
     else: raise CookieAuthRejected("Socket failed (%s)" % exc, cookie_path)
 
-def get_protocolinfo(control_socket):
+def get_protocolinfo(controller):
   """
   Issues a PROTOCOLINFO query to a control socket, getting information about
   the tor process running on it. If the socket is already closed then it is
   first reconnected.
   
   Arguments:
-    control_socket (stem.socket.ControlSocket) - connected tor control socket
+    controller (stem.socket.ControlSocket or stem.control.BaseController) -
+      tor controller connection
   
   Returns:
     stem.connection.ProtocolInfoResponse provided by tor
@@ -642,8 +645,7 @@ def get_protocolinfo(control_socket):
   """
   
   try:
-    control_socket.send("PROTOCOLINFO 1")
-    protocolinfo_response = control_socket.recv()
+    protocolinfo_response = _msg(controller, "PROTOCOLINFO 1")
   except:
     protocolinfo_response = None
   
@@ -651,17 +653,22 @@ def get_protocolinfo(control_socket):
   # next followed by authentication. Transparently reconnect if that happens.
   
   if not protocolinfo_response or str(protocolinfo_response) == "Authentication required.":
-    control_socket.connect()
+    controller.connect()
     
     try:
-      control_socket.send("PROTOCOLINFO 1")
-      protocolinfo_response = control_socket.recv()
+      protocolinfo_response = _msg(controller, "PROTOCOLINFO 1")
     except stem.socket.SocketClosed, exc:
       raise stem.socket.SocketError(exc)
   
   ProtocolInfoResponse.convert(protocolinfo_response)
   
-  # attempt ot expand relative cookie paths via the control port or socket file
+  # attempt to expand relative cookie paths via the control port or socket file
+  
+  if isinstance(controller, stem.socket.ControlSocket):
+    control_socket = controller
+  else:
+    control_socket = controller.get_socket()
+  
   if isinstance(control_socket, stem.socket.ControlPort):
     if control_socket.get_address() == "127.0.0.1":
       pid_method = stem.util.system.get_pid_by_port
@@ -672,6 +679,17 @@ def get_protocolinfo(control_socket):
   
   return protocolinfo_response
 
+def _msg(controller, message):
+  """
+  Sends and receives a message with either a ControlSocket or BaseController.
+  """
+  
+  if isinstance(controller, stem.socket.ControlSocket):
+    controller.send(message)
+    return controller.recv()
+  else:
+    return controller.msg(message)
+
 def _expand_cookie_path(protocolinfo_response, pid_resolver, pid_resolution_arg):
   """
   Attempts to expand a relative cookie path with the given pid resolver. This
diff --git a/test/integ/connection/authentication.py b/test/integ/connection/authentication.py
index f80b848..c209d2b 100644
--- a/test/integ/connection/authentication.py
+++ b/test/integ/connection/authentication.py
@@ -80,7 +80,7 @@ class TestAuthenticate(unittest.TestCase):
   def setUp(self):
     test.runner.require_control(self)
   
-  def test_authenticate_general(self):
+  def test_authenticate_general_socket(self):
     """
     Tests that the authenticate function can authenticate to our socket.
     """
@@ -88,7 +88,17 @@ class TestAuthenticate(unittest.TestCase):
     runner = test.runner.get_runner()
     with runner.get_tor_socket(False) as control_socket:
       stem.connection.authenticate(control_socket, test.runner.CONTROL_PASSWORD, runner.get_chroot())
-      test.runner.exercise_socket(self, control_socket)
+      test.runner.exercise_controller(self, control_socket)
+  
+  def test_authenticate_general_controller(self):
+    """
+    Tests that the authenticate function can authenticate via a Controller.
+    """
+    
+    runner = test.runner.get_runner()
+    with runner.get_tor_controller(False) as controller:
+      stem.connection.authenticate(controller, test.runner.CONTROL_PASSWORD, runner.get_chroot())
+      test.runner.exercise_controller(self, controller)
   
   def test_authenticate_general_example(self):
     """
@@ -108,7 +118,7 @@ class TestAuthenticate(unittest.TestCase):
     try:
       # this authenticate call should work for everything but password-only auth
       stem.connection.authenticate(control_socket, chroot_path = runner.get_chroot())
-      test.runner.exercise_socket(self, control_socket)
+      test.runner.exercise_controller(self, control_socket)
     except stem.connection.IncorrectSocketType:
       self.fail()
     except stem.connection.MissingPassword:
@@ -117,7 +127,7 @@ class TestAuthenticate(unittest.TestCase):
       
       try:
         stem.connection.authenticate_password(control_socket, controller_password)
-        test.runner.exercise_socket(self, control_socket)
+        test.runner.exercise_controller(self, control_socket)
       except stem.connection.PasswordAuthFailed:
         self.fail()
     except stem.connection.AuthenticationFailure:
@@ -144,7 +154,7 @@ class TestAuthenticate(unittest.TestCase):
         self.assertRaises(stem.connection.MissingPassword, stem.connection.authenticate, control_socket)
       else:
         stem.connection.authenticate(control_socket, chroot_path = runner.get_chroot())
-        test.runner.exercise_socket(self, control_socket)
+        test.runner.exercise_controller(self, control_socket)
     
     # tests with the incorrect password
     with runner.get_tor_socket(False) as control_socket:
@@ -152,12 +162,12 @@ class TestAuthenticate(unittest.TestCase):
         self.assertRaises(stem.connection.IncorrectPassword, stem.connection.authenticate, control_socket, "blarg")
       else:
         stem.connection.authenticate(control_socket, "blarg", runner.get_chroot())
-        test.runner.exercise_socket(self, control_socket)
+        test.runner.exercise_controller(self, control_socket)
     
     # tests with the right password
     with runner.get_tor_socket(False) as control_socket:
       stem.connection.authenticate(control_socket, test.runner.CONTROL_PASSWORD, runner.get_chroot())
-      test.runner.exercise_socket(self, control_socket)
+      test.runner.exercise_controller(self, control_socket)
   
   def test_authenticate_none(self):
     """
@@ -299,7 +309,7 @@ class TestAuthenticate(unittest.TestCase):
         elif auth_type == stem.connection.AuthMethod.COOKIE:
           stem.connection.authenticate_cookie(control_socket, auth_arg)
         
-        test.runner.exercise_socket(self, control_socket)
+        test.runner.exercise_controller(self, control_socket)
       except stem.connection.AuthenticationFailure, exc:
         # authentication functions should re-attach on failure
         self.assertTrue(control_socket.is_alive())
diff --git a/test/integ/connection/connect.py b/test/integ/connection/connect.py
index 29cf2dc..ac66ceb 100644
--- a/test/integ/connection/connect.py
+++ b/test/integ/connection/connect.py
@@ -34,7 +34,7 @@ class TestConnect(unittest.TestCase):
       controller = stem.connection.Controller.NONE)
     
     if test.runner.Torrc.PORT in runner.get_options():
-      test.runner.exercise_socket(self, control_socket)
+      test.runner.exercise_controller(self, control_socket)
       control_socket.close()
     else:
       self.assertEquals(control_socket, None)
@@ -53,7 +53,7 @@ class TestConnect(unittest.TestCase):
       controller = stem.connection.Controller.NONE)
     
     if test.runner.Torrc.SOCKET in runner.get_options():
-      test.runner.exercise_socket(self, control_socket)
+      test.runner.exercise_controller(self, control_socket)
       control_socket.close()
     else:
       self.assertEquals(control_socket, None)
diff --git a/test/runner.py b/test/runner.py
index 6702b8b..66636f2 100644
--- a/test/runner.py
+++ b/test/runner.py
@@ -8,7 +8,7 @@ TorInaccessable - Tor can't be queried for the information
 
 require_control - skips the test unless tor provides a controller endpoint
 require_version - skips the test unless we meet a tor version requirement
-exercise_socket - Does a basic sanity check that a control socket can be used
+exercise_controller - basic sanity check that a controller connection can be used
 
 get_runner - Singleton for fetching our runtime context.
 Runner - Runtime context for our integration tests.
@@ -25,7 +25,8 @@ Runner - Runtime context for our integration tests.
   |- get_tor_cwd - current working directory of our tor process
   |- get_chroot - provides the path of our emulated chroot if we have one
   |- get_pid - process id of our tor process
-  |- get_tor_socket - provides a socket to the tor instance
+  |- get_tor_socket - provides a socket to our test instance
+  |- get_tor_controller - provides a controller for our test instance
   |- get_tor_version - provides the version of tor we're running against
   +- get_tor_command - provides the command used to start tor
 """
@@ -115,14 +116,15 @@ def require_version(test_case, req_version):
   if get_runner().get_tor_version() < req_version:
     test_case.skipTest("(requires %s)" % req_version)
 
-def exercise_socket(test_case, control_socket):
+def exercise_controller(test_case, controller):
   """
   Checks that we can now use the socket by issuing a 'GETINFO config-file'
   query.
   
   Arguments:
     test_case (unittest.TestCase) - test being ran
-    control_socket (stem.socket.ControlSocket) - socket to be tested
+    controller (stem.socket.ControlSocket or stem.control.BaseController) -
+      tor controller connection to be authenticated
   """
   
   runner = get_runner()
@@ -131,8 +133,12 @@ def exercise_socket(test_case, control_socket):
   if chroot_path and torrc_path.startswith(chroot_path):
     torrc_path = torrc_path[len(chroot_path):]
   
-  control_socket.send("GETINFO config-file")
-  config_file_response = control_socket.recv()
+  if isinstance(controller, stem.socket.ControlSocket):
+    controller.send("GETINFO config-file")
+    config_file_response = controller.recv()
+  else:
+    config_file_response = controller.msg("GETINFO config-file")
+  
   test_case.assertEquals("config-file=%s\nOK" % torrc_path, str(config_file_response))
 
 def get_runner():
@@ -429,7 +435,7 @@ class Runner:
   
   def get_tor_socket(self, authenticate = True):
     """
-    Provides a socket connected to the tor test instance's control socket.
+    Provides a socket connected to our tor test instance.
     
     Arguments:
       authenticate (bool) - if True then the socket is authenticated
@@ -452,6 +458,24 @@ class Runner:
     
     return control_socket
   
+  def get_tor_controller(self, authenticate = True):
+    """
+    Provides a controller connected to our tor test instance.
+    
+    Arguments:
+      authenticate (bool) - if True then the socket is authenticated
+    
+    Returns:
+      stem.socket.BaseController connected with our testing instance
+    
+    Raises:
+      TorInaccessable if tor can't be connected to
+    """
+    
+    # TODO: replace with our general controller when we have one
+    control_socket = self.get_tor_socket(authenticate)
+    return stem.control.BaseController(control_socket)
+  
   def get_tor_version(self):
     """
     Queries our test instance for tor's version.



More information about the tor-commits mailing list