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):