commit bca28c1e409a318bf9303a635530382c13c17115 Author: Damian Johnson <atagar@torproject.org> Date: Thu May 17 09:29:04 2012 -0700 Replacing enums with class instances for connect_* functions The connect_port and connect_socket_file convenience functions accept a controller type that we would like to return. I had been using an enum for this argument, but this means that we both (a) need to manually update it when adding controllers and (b) can't recognize controller the user makes. Replacing this with a class instance, like what TorCtl does. Tested by running against a script for printing BW events. --- stem/connection.py | 42 +++++++++++++++++++------------------ test/integ/connection/connect.py | 4 +- 2 files changed, 24 insertions(+), 22 deletions(-) diff --git a/stem/connection.py b/stem/connection.py index 825edf5..b81a8f7 100644 --- a/stem/connection.py +++ b/stem/connection.py @@ -90,9 +90,6 @@ import stem.util.enum import stem.util.system import stem.util.log as log -# enums representing classes that the connect_* methods can return -Controller = stem.util.enum.Enum("NONE") - # Methods by which a controller can authenticate to the control port. Tor gives # a list of all the authentication methods it will accept in response to # PROTOCOLINFO queries. @@ -194,7 +191,7 @@ AUTHENTICATE_EXCEPTIONS = ( AuthenticationFailure, ) -def connect_port(control_addr = "127.0.0.1", control_port = 9051, password = None, chroot_path = None, controller = Controller.NONE): +def connect_port(control_addr = "127.0.0.1", control_port = 9051, password = None, chroot_path = None, controller = None): """ Convenience function for quickly getting a control connection. This is very handy for debugging or CLI setup, handling setup and prompting for a password @@ -202,15 +199,16 @@ def connect_port(control_addr = "127.0.0.1", control_port = 9051, password = Non description of the problem and returns None. Arguments: - control_addr (str) - ip address of the controller - control_port (int) - port number of the controller - password (str) - passphrase to authenticate to the socket - chroot_path (str) - path prefix if in a chroot environment - controller (Controller) - controller type to be returned + control_addr (str) - ip address of the controller + control_port (int) - port number of the controller + password (str) - passphrase to authenticate to the socket + chroot_path (str) - path prefix if in a chroot environment + controller (Class) - BaseController subclass to be returned, this provides + a ControlSocket if None Returns: - Authenticated control connection, the type based on the controller enum... - Controller.NONE => stem.socket.ControlPort + Authenticated control connection, the type based on the controller + argument. """ # TODO: replace the controller arg's default when we have something better @@ -223,16 +221,17 @@ def connect_port(control_addr = "127.0.0.1", control_port = 9051, password = Non return _connect(control_port, password, chroot_path, controller) -def connect_socket_file(socket_path = "/var/run/tor/control", password = None, chroot_path = None, controller = Controller.NONE): +def connect_socket_file(socket_path = "/var/run/tor/control", password = None, chroot_path = None, controller = None): """ Convenience function for quickly getting a control connection. For more information see the connect_port function. Arguments: - socket_path (str) - path where the control socket is located - password (str) - passphrase to authenticate to the socket - chroot_path (str) - path prefix if in a chroot environment - controller (Controller) - controller type to be returned + socket_path (str) - path where the control socket is located + password (str) - passphrase to authenticate to the socket + chroot_path (str) - path prefix if in a chroot environment + controller (Class) - BaseController subclass to be returned, this provides + a ControlSocket if None Returns: Authenticated control connection, the type based on the controller enum. @@ -252,9 +251,10 @@ def _connect(control_socket, password, chroot_path, controller): Arguments: control_socket (stem.socket.ControlSocket) - socket being authenticated to - password (str) - passphrase to authenticate to the socket - chroot_path (str) - path prefix if in a chroot environment - controller (Controller) - controller type to be returned + password (str) - passphrase to authenticate to the socket + chroot_path (str) - path prefix if in a chroot environment + controller (Class) - BaseController subclass to be returned, this provides + a ControlSocket if None Returns: Authenticated control connection with a type based on the controller enum. @@ -263,8 +263,10 @@ def _connect(control_socket, password, chroot_path, controller): try: authenticate(control_socket, password, chroot_path) - if controller == Controller.NONE: + if controller == None: return control_socket + else: + return controller(control_socket) except MissingPassword: assert password is None, "BUG: authenticate raised MissingPassword despite getting one" diff --git a/test/integ/connection/connect.py b/test/integ/connection/connect.py index ac66ceb..3af249b 100644 --- a/test/integ/connection/connect.py +++ b/test/integ/connection/connect.py @@ -31,7 +31,7 @@ class TestConnect(unittest.TestCase): control_port = test.runner.CONTROL_PORT, password = test.runner.CONTROL_PASSWORD, chroot_path = runner.get_chroot(), - controller = stem.connection.Controller.NONE) + controller = None) if test.runner.Torrc.PORT in runner.get_options(): test.runner.exercise_controller(self, control_socket) @@ -50,7 +50,7 @@ class TestConnect(unittest.TestCase): socket_path = test.runner.CONTROL_SOCKET_PATH, password = test.runner.CONTROL_PASSWORD, chroot_path = runner.get_chroot(), - controller = stem.connection.Controller.NONE) + controller = None) if test.runner.Torrc.SOCKET in runner.get_options(): test.runner.exercise_controller(self, control_socket)