[tor-commits] [stem/master] Making connect()'s password prompt opt-in

[atagar at torproject.org]

commit be37a4f8a7c388970518e1502ef6d6194743a9d1
Author: Damian Johnson <atagar at torproject.org>
Date:   Mon Apr 7 09:58:02 2014 -0700

    Making connect()'s password prompt opt-in
    
    Making use of stdin opt-in allows the connect() function to be used more safely
    in quick and dirty scripts. Users can then allow a password prompt if waiting
    on stdin is safe.
---
 stem/connection.py              |   30 ++++++++++++++++++++----------
 test/unit/connection/connect.py |    8 ++++----
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/stem/connection.py b/stem/connection.py
index 35524d7..c49ff85 100644
--- a/stem/connection.py
+++ b/stem/connection.py
@@ -186,6 +186,7 @@ CONNECT_MESSAGES = {
   'incorrect_password': "Incorrect password",
   'no_control_port': "Unable to connect to tor. Maybe it's running without a ControlPort?",
   'password_prompt': "Tor controller password:",
+  'needs_password': "Tor requires a password to authenticate",
   'socket_doesnt_exist': "The socket file you specified ({path}) doesn't exist",
   'tor_isnt_running': "Unable to connect to tor. Are you sure it's running?",
   'unable_to_use_port': "Unable to connect to {address}:{port}: {error}",
@@ -198,7 +199,7 @@ CONNECT_MESSAGES = {
 }
 
 
-def connect(control_port = ('127.0.0.1', 9051), control_socket = '/var/run/tor/control', password = None, chroot_path = None, controller = stem.control.Controller):
+def connect(control_port = ('127.0.0.1', 9051), control_socket = '/var/run/tor/control', password = None, password_prompt = False, chroot_path = None, controller = stem.control.Controller):
   """
   Convenience function for quickly getting a control connection. This is very
   handy for debugging or CLI setup, handling setup and prompting for a password
@@ -216,6 +217,8 @@ def connect(control_port = ('127.0.0.1', 9051), control_socket = '/var/run/tor/c
   :param tuple contol_port: address and port tuple, for instance **('127.0.0.1', 9051)**
   :param str path: path where the control socket is located
   :param str password: passphrase to authenticate to the socket
+  :param bool password_prompt: prompt for the controller password if it wasn't
+    supplied
   :param str chroot_path: path prefix if in a chroot environment
   :param Class controller: :class:`~stem.control.BaseController` subclass to be
     returned, this provides a :class:`~stem.socket.ControlSocket` if **None**
@@ -271,7 +274,7 @@ def connect(control_port = ('127.0.0.1', 9051), control_socket = '/var/run/tor/c
     print error_msg
     return None
 
-  return _connect_auth(control_connection, password, chroot_path, controller)
+  return _connect_auth(control_connection, password, password_prompt, chroot_path, controller)
 
 
 def connect_port(address = "127.0.0.1", port = 9051, password = None, chroot_path = None, controller = stem.control.Controller):
@@ -300,7 +303,7 @@ def connect_port(address = "127.0.0.1", port = 9051, password = None, chroot_pat
     print exc
     return None
 
-  return _connect_auth(control_port, password, chroot_path, controller)
+  return _connect_auth(control_port, password, True, chroot_path, controller)
 
 
 def connect_socket_file(path = "/var/run/tor/control", password = None, chroot_path = None, controller = stem.control.Controller):
@@ -330,16 +333,18 @@ def connect_socket_file(path = "/var/run/tor/control", password = None, chroot_p
     print exc
     return None
 
-  return _connect_auth(control_socket, password, chroot_path, controller)
+  return _connect_auth(control_socket, password, True, chroot_path, controller)
 
 
-def _connect_auth(control_socket, password, chroot_path, controller):
+def _connect_auth(control_socket, password, password_prompt, chroot_path, controller):
   """
   Helper for the connect_* functions that authenticates the socket and
   constructs the controller.
 
   :param stem.socket.ControlSocket control_socket: socket being authenticated to
   :param str password: passphrase to authenticate to the socket
+  :param bool password_prompt: prompt for the controller password if it wasn't
+    supplied
   :param str chroot_path: path prefix if in a chroot environment
   :param Class controller: :class:`~stem.control.BaseController` subclass to be
     returned, this provides a :class:`~stem.socket.ControlSocket` if **None**
@@ -375,13 +380,18 @@ def _connect_auth(control_socket, password, chroot_path, controller):
       control_socket.close()
       raise ValueError(CONNECT_MESSAGES['missing_password_bug'])
 
-    try:
-      password = getpass.getpass(CONNECT_MESSAGES['password_prompt'] + ' ')
-    except KeyboardInterrupt:
+    if password_prompt:
+      try:
+        password = getpass.getpass(CONNECT_MESSAGES['password_prompt'] + ' ')
+      except KeyboardInterrupt:
+        control_socket.close()
+        return None
+
+      return _connect_auth(control_socket, password, password_prompt, chroot_path, controller)
+    else:
+      print CONNECT_MESSAGES['needs_password']
       control_socket.close()
       return None
-
-    return _connect_auth(control_socket, password, chroot_path, controller)
   except UnreadableCookieFile as exc:
     print CONNECT_MESSAGES['unreadable_cookie_file'].format(path = exc.cookie_path, issue = str(exc))
     control_socket.close()
diff --git a/test/unit/connection/connect.py b/test/unit/connection/connect.py
index 52039f2..9265a5a 100644
--- a/test/unit/connection/connect.py
+++ b/test/unit/connection/connect.py
@@ -77,11 +77,11 @@ class TestConnect(unittest.TestCase):
   def test_auth_success(self, authenticate_mock):
     control_socket = Mock()
 
-    stem.connection._connect_auth(control_socket, None, None, None)
+    stem.connection._connect_auth(control_socket, None, False, None, None)
     authenticate_mock.assert_called_with(control_socket, None, None)
     authenticate_mock.reset_mock()
 
-    stem.connection._connect_auth(control_socket, 's3krit!!!', '/my/chroot', None)
+    stem.connection._connect_auth(control_socket, 's3krit!!!', False, '/my/chroot', None)
     authenticate_mock.assert_called_with(control_socket, 's3krit!!!', '/my/chroot')
 
   @patch('getpass.getpass')
@@ -100,7 +100,7 @@ class TestConnect(unittest.TestCase):
     authenticate_mock.side_effect = authenticate_mock_func
     getpass_mock.return_value = 'my_password'
 
-    stem.connection._connect_auth(control_socket, None, None, None)
+    stem.connection._connect_auth(control_socket, None, True, None, None)
     authenticate_mock.assert_any_call(control_socket, None, None)
     authenticate_mock.assert_any_call(control_socket, 'my_password', None)
 
@@ -129,7 +129,7 @@ class TestConnect(unittest.TestCase):
     self._assert_authenticate_fails_with(control_socket, stdout_mock, 'Unable to authenticate: crazy failure')
 
   def _assert_authenticate_fails_with(self, control_socket, stdout_mock, msg):
-    result = stem.connection._connect_auth(control_socket, None, None, None)
+    result = stem.connection._connect_auth(control_socket, None, False, None, None)
 
     if result is not None:
       self.fail()  # _connect_auth() was successful