webssh

commit 33703f0ca35c033bde0a603a4ccbcd741e84e8b9
parent c873b50522c3ce97107cd83efdd79639569b7e86
Author: Sheng <webmaster0115@gmail.com>
Date:   Mon,  8 Oct 2018 22:37:03 +0800

Updated lookup_hostname

Diffstat:
M
tests/test_app.py
 | 
3
++-
M
webssh/handler.py
 | 
12
+++++++-----

2 files changed, 9 insertions(+), 6 deletions(-)
diff --git a/tests/test_app.py b/tests/test_app.py
@@ -536,4 +536,5 @@ class TestAppWithRejectPolicy(OtherTestBase):
         data = json.loads(to_str(response.body))
         self.assertIsNone(data['id'])
         self.assertIsNone(data['encoding'])
-        self.assertEqual('Connection to 127.0.0.1 is not allowed.', data['status'])  # noqa
+        message = 'Connection to {}:{} is not allowed.'.format(self.body['hostname'], self.sshserver_port) # noqa
+        self.assertEqual(message, data['status'])
diff --git a/webssh/handler.py b/webssh/handler.py
@@ -152,12 +152,14 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
             raise InvalidValueError('Invalid hostname: {}'.format(value))
         return value
 
-    def lookup_hostname(self, hostname):
+    def lookup_hostname(self, hostname, port):
         if isinstance(self.policy, paramiko.RejectPolicy):
-            if self.ssh_client._system_host_keys.lookup(hostname) is None:
-                if self.ssh_client._host_keys.lookup(hostname) is None:
+            key = hostname if port == 22 else '[{}]:{}'.format(hostname, port)
+            if self.ssh_client._system_host_keys.lookup(key) is None:
+                if self.ssh_client._host_keys.lookup(key) is None:
                     raise ValueError(
-                        'Connection to {} is not allowed.'.format(hostname)
+                        'Connection to {}:{} is not allowed.'.format(
+                            hostname, port)
                     )
 
     def get_port(self):
@@ -174,8 +176,8 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
 
     def get_args(self):
         hostname = self.get_hostname()
-        self.lookup_hostname(hostname)
         port = self.get_port()
+        self.lookup_hostname(hostname, port)
         username = self.get_value('username')
         password = self.get_argument('password', u'')
         privatekey = self.get_privatekey()