Refactored handler.py

tests/test_app.py

@@ -40,12 +40,12 @@ class TestAppBase(AsyncHTTPTestCase):
             self.assertEqual(response.code, 400)
             self.assertIn(b'Bad Request', response.body)
 
-    def assert_status_in(self, data, status):
+    def assert_status_in(self, status, data):
         self.assertIsNone(data['encoding'])
         self.assertIsNone(data['id'])
         self.assertIn(status, data['status'])
 
-    def assert_status_equal(self, data, status):
+    def assert_status_equal(self, status, data):
         self.assertIsNone(data['encoding'])
         self.assertIsNone(data['id'])
         self.assertEqual(status, data['status'])
@@ -172,7 +172,7 @@ class TestAppBasic(TestAppBase):
 
     def test_app_with_wrong_credentials(self):
         response = self.sync_post('/', self.body + 's')
-        self.assert_status_in(json.loads(to_str(response.body)), 'Authentication failed.') # noqa
+        self.assert_status_in('Authentication failed.', json.loads(to_str(response.body))) # noqa
 
     def test_app_with_correct_credentials(self):
         response = self.sync_post('/', self.body)
@@ -442,10 +442,10 @@ class TestAppBasic(TestAppBase):
         self.body_dict.update(username='keyonly', password='foo')
         response = yield self.async_post('/', self.body_dict)
         self.assertEqual(response.code, 200)
-        self.assert_status_in(json.loads(to_str(response.body)), 'Bad authentication type') # noqa
+        self.assert_status_in('Bad authentication type', json.loads(to_str(response.body))) # noqa
 
     @tornado.testing.gen_test
-    def test_app_with_user_pass2fa_with_correct_password_and_passcode(self):
+    def test_app_with_user_pass2fa_with_correct_passwords(self):
         self.body_dict.update(username='pass2fa', password='password',
                               totp='passcode')
         response = yield self.async_post('/', self.body_dict)
@@ -454,25 +454,7 @@ class TestAppBasic(TestAppBase):
         self.assert_status_none(data)
 
     @tornado.testing.gen_test
-    def test_app_with_user_pass2fa_with_wrong_password(self):
+    def test_app_with_user_pass2fa_with_wrong_pkey_correct_passwords(self):
-        self.body_dict.update(username='pass2fa', password='wrongpassword',
-                              totp='passcode')
-        response = yield self.async_post('/', self.body_dict)
-        self.assertEqual(response.code, 200)
-        data = json.loads(to_str(response.body))
-        self.assertIn('Authentication failed', data['status'])
-
-    @tornado.testing.gen_test
-    def test_app_with_user_pass2fa_with_wrong_passcode(self):
-        self.body_dict.update(username='pass2fa', password='password',
-                              totp='wrongpasscode')
-        response = yield self.async_post('/', self.body_dict)
-        self.assertEqual(response.code, 200)
-        data = json.loads(to_str(response.body))
-        self.assertIn('Authentication failed', data['status'])
-
-    @tornado.testing.gen_test
-    def test_app_with_user_pass2fa_with_wrong_pkey_correct_passwords(self):  # noqa
         url = self.get_url('/')
         privatekey = read_file(make_tests_data_path('user_rsa_key'))
         self.body_dict.update(username='pass2fa', password='password',
@@ -482,7 +464,7 @@ class TestAppBasic(TestAppBase):
         self.assert_status_none(data)
 
     @tornado.testing.gen_test
-    def test_app_with_user_pkey2fa_with_correct_password_and_passcode(self):
+    def test_app_with_user_pkey2fa_with_correct_passwords(self):
         url = self.get_url('/')
         privatekey = read_file(make_tests_data_path('user_rsa_key'))
         self.body_dict.update(username='pkey2fa', password='password',
@@ -499,7 +481,7 @@ class TestAppBasic(TestAppBase):
                               privatekey=privatekey, totp='passcode')
         response = yield self.async_post(url, self.body_dict)
         data = json.loads(to_str(response.body))
-        self.assertIn('Authentication failed', data['status'])
+        self.assert_status_in('Authentication failed', data)
 
     @tornado.testing.gen_test
     def test_app_with_user_pkey2fa_with_wrong_passcode(self):
@@ -509,7 +491,17 @@ class TestAppBasic(TestAppBase):
                               privatekey=privatekey, totp='wrongpasscode')
         response = yield self.async_post(url, self.body_dict)
         data = json.loads(to_str(response.body))
-        self.assertIn('Authentication failed', data['status'])
+        self.assert_status_in('Authentication failed', data)
+
+    @tornado.testing.gen_test
+    def test_app_with_user_pkey2fa_with_empty_passcode(self):
+        url = self.get_url('/')
+        privatekey = read_file(make_tests_data_path('user_rsa_key'))
+        self.body_dict.update(username='pkey2fa', password='password',
+                              privatekey=privatekey, totp='')
+        response = yield self.async_post(url, self.body_dict)
+        data = json.loads(to_str(response.body))
+        self.assert_status_in('Need a verification code', data)
 
 
 class OtherTestBase(TestAppBase):
@@ -747,13 +739,13 @@ class TestAppWithCrossOriginOperation(OtherTestBase):
     def test_app_with_wrong_event_origin(self):
         body = dict(self.body, _origin='localhost')
         response = yield self.async_post('/', body)
-        self.assert_status_equal(json.loads(to_str(response.body)), 'Cross origin operation is not allowed.') # noqa
+        self.assert_status_equal('Cross origin operation is not allowed.', json.loads(to_str(response.body))) # noqa
 
     @tornado.testing.gen_test
     def test_app_with_wrong_header_origin(self):
         headers = dict(Origin='localhost')
         response = yield self.async_post('/', self.body, headers=headers)
-        self.assert_status_equal(json.loads(to_str(response.body)), 'Cross origin operation is not allowed.') # noqa
+        self.assert_status_equal('Cross origin operation is not allowed.', json.loads(to_str(response.body)), ) # noqa
 
     @tornado.testing.gen_test
     def test_app_with_correct_event_origin(self):

webssh/handler.py

@@ -36,80 +36,68 @@ swallow_http_errors = True
 redirecting = None
 
 
-def make_handler(password, totp):
+class InvalidValueError(Exception):
+    pass
 
-    def handler(title, instructions, prompt_list):
+
+class SSHClient(paramiko.SSHClient):
+
+    def handler(self, title, instructions, prompt_list):
         answers = []
         for prompt_, _ in prompt_list:
             prompt = prompt_.strip().lower()
             if prompt.startswith('password'):
-                answers.append(password)
+                answers.append(self.password)
             elif prompt.startswith('verification'):
-                answers.append(totp)
+                answers.append(self.totp)
             else:
                 raise ValueError('Unknown prompt: {}'.format(prompt_))
         return answers
 
-    return handler
+    def auth_interactive(self, username, handler):
+        if not self.totp:
+            raise ValueError('Need a verification code for 2fa.')
+        self._transport.auth_interactive(username, handler)
 
+    def _auth(self, username, password, pkey, *args):
+        self.password = password
+        saved_exception = None
+        two_factor = False
+        allowed_types = set()
+        two_factor_types = {'keyboard-interactive', 'password'}
 
-def auth_interactive(transport, username, handler):
+        if pkey is not None:
-    if not handler:
+            logging.info('Trying publickey authentication')
-        raise ValueError('Need a verification code for 2fa.')
+            try:
-    transport.auth_interactive(username, handler)
+                allowed_types = set(
+                    self._transport.auth_publickey(username, pkey)
+                )
+                two_factor = allowed_types & two_factor_types
+                if not two_factor:
+                    return
+            except paramiko.SSHException as e:
+                saved_exception = e
 
+        if two_factor:
+            logging.info('Trying publickey 2fa')
+            return self.auth_interactive(username, self.handler)
 
-def auth(self, username, password, pkey, *args):
+        if password is not None:
-    handler = None
+            logging.info('Trying password authentication')
-    saved_exception = None
+            try:
-    two_factor = False
+                self._transport.auth_password(username, password)
-    allowed_types = set()
-    two_factor_types = {"keyboard-interactive", "password"}
-
-    if self._totp:
-        handler = make_handler(password, self._totp)
-
-    if pkey is not None:
-        logging.info('Trying public key authentication')
-        try:
-            allowed_types = set(
-                self._transport.auth_publickey(username, pkey)
-            )
-            two_factor = allowed_types & two_factor_types
-            if not two_factor:
                 return
-        except paramiko.SSHException as e:
+            except paramiko.SSHException as e:
-            saved_exception = e
+                saved_exception = e
+                allowed_types = set(getattr(e, 'allowed_types', []))
+                two_factor = allowed_types & two_factor_types
 
-    if two_factor:
+        if two_factor:
-        logging.info('Trying publickey 2fa')
+            logging.info('Trying password 2fa')
-        return auth_interactive(self._transport, username, handler)
+            return self.auth_interactive(username, self.handler)
 
-    if password is not None:
+        assert saved_exception is not None
-        logging.info('Trying password authentication')
-        try:
-            self._transport.auth_password(username, password)
-            return
-        except paramiko.SSHException as e:
-            saved_exception = e
-            allowed_types = set(getattr(e, 'allowed_types', []))
-            two_factor = allowed_types & two_factor_types
-
-    if two_factor:
-        logging.info('Trying password 2fa')
-        return auth_interactive(self._transport, username, handler)
-
-    # if we got an auth-failed exception earlier, re-raise it
-    if saved_exception is not None:
         raise saved_exception
-    raise paramiko.SSHException("No authentication methods available")
-
-
-paramiko.client.SSHClient._auth = auth
-
-
-class InvalidValueError(Exception):
-    pass
 
 
 class PrivateKey(object):
@@ -327,7 +315,7 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
             super(IndexHandler, self).write_error(status_code, **kwargs)
 
     def get_ssh_client(self):
-        ssh = paramiko.SSHClient()
+        ssh = SSHClient()
         ssh._system_host_keys = self.host_keys_settings['system_host_keys']
         ssh._host_keys = self.host_keys_settings['host_keys']
         ssh._host_keys_filename = self.host_keys_settings['host_keys_filename']
@@ -392,7 +380,7 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
         else:
             pkey = None
 
-        self.ssh_client._totp = totp
+        self.ssh_client.totp = totp
         args = (hostname, port, username, password, pkey)
         logging.debug(args)