Support 2fa

tests/sshserver.py

@@ -57,6 +57,8 @@ class Server(paramiko.ServerInterface):
         self.shell_event = threading.Event()
         self.exec_event = threading.Event()
         self.encoding = random.choice(self.encodings)
+        self.password_verified = False
+        self.key_verified = False
 
     def check_channel_request(self, kind, chanid):
         if kind == 'session':
@@ -73,11 +75,50 @@ class Server(paramiko.ServerInterface):
         print('Auth attempt with username: {!r} & key: {!r}'.format(username, u(hexlify(key.get_fingerprint())))) # noqa
         if (username in ['robey', 'keyonly']) and (key == self.good_pub_key):
             return paramiko.AUTH_SUCCESSFUL
+        if username == 'pkey2fa' and key == self.good_pub_key:
+            self.key_verified = True
+            return paramiko.AUTH_PARTIALLY_SUCCESSFUL
         return paramiko.AUTH_FAILED
 
+    def check_auth_interactive(self, username, submethods):
+        if username in ['pass2fa', 'pkey2fa']:
+            self.username = username
+            prompt = 'Verification code: ' if self.password_verified else 'Password: '  # noqa
+            print(username, prompt)
+            return paramiko.InteractiveQuery('', '', prompt)
+        return paramiko.AUTH_FAILED
+
+    def check_auth_interactive_response(self, responses):
+        if self.username in ['pass2fa', 'pkey2fa']:
+            if not self.password_verified:
+                if responses[0] == 'password':
+                    print('password verified')
+                    self.password_verified = True
+                    if self.username == 'pkey2fa':
+                        return self.check_auth_interactive(self.username, '')
+                else:
+                    print('wrong password: {}'.format(responses[0]))
+                    return paramiko.AUTH_FAILED
+            else:
+                if responses[0] == 'passcode':
+                    print('totp verified')
+                    return paramiko.AUTH_SUCCESSFUL
+                else:
+                    print('wrong totp: {}'.format(responses[0]))
+                    return paramiko.AUTH_FAILED
+        else:
+            return paramiko.AUTH_FAILED
+
     def get_allowed_auths(self, username):
         if username == 'keyonly':
             return 'publickey'
+        if username == 'pass2fa':
+            return 'keyboard-interactive'
+        if username == 'pkey2fa':
+            if not self.key_verified:
+                return 'publickey'
+            else:
+                return 'keyboard-interactive'
         return 'password,publickey'
 
     def check_channel_exec_request(self, channel, command):

tests/test_app.py

@@ -444,6 +444,73 @@ class TestAppBasic(TestAppBase):
         self.assertEqual(response.code, 200)
         self.assert_status_in(json.loads(to_str(response.body)), 'Bad authentication type') # noqa
 
+    @tornado.testing.gen_test
+    def test_app_with_user_pass2fa_with_correct_password_and_passcode(self):
+        self.body_dict.update(username='pass2fa', password='password',
+                              totp='passcode')
+        response = yield self.async_post('/', self.body_dict)
+        self.assertEqual(response.code, 200)
+        data = json.loads(to_str(response.body))
+        self.assert_status_none(data)
+
+    @tornado.testing.gen_test
+    def test_app_with_user_pass2fa_with_wrong_password(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',
+                              privatekey=privatekey, totp='passcode')
+        response = yield self.async_post(url, self.body_dict)
+        data = json.loads(to_str(response.body))
+        self.assert_status_none(data)
+
+    @tornado.testing.gen_test
+    def test_app_with_user_pkey2fa_with_correct_password_and_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='passcode')
+        response = yield self.async_post(url, self.body_dict)
+        data = json.loads(to_str(response.body))
+        self.assert_status_none(data)
+
+    @tornado.testing.gen_test
+    def test_app_with_user_pkey2fa_with_wrong_password(self):
+        url = self.get_url('/')
+        privatekey = read_file(make_tests_data_path('user_rsa_key'))
+        self.body_dict.update(username='pkey2fa', password='wrongpassword',
+                              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'])
+
+    @tornado.testing.gen_test
+    def test_app_with_user_pkey2fa_with_wrong_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='wrongpasscode')
+        response = yield self.async_post(url, self.body_dict)
+        data = json.loads(to_str(response.body))
+        self.assertIn('Authentication failed', data['status'])
+
 
 class OtherTestBase(TestAppBase):
     sshserver_port = 3300