webssh

commit fc30ead69e63c335513b82de1f03defdcab4c1a0
parent 049baad90948ce9774ca55ef208a79f06806926c
Author: Sheng <webmaster0115@gmail.com>
Date:   Wed,  9 Oct 2019 11:40:07 +0800

Refactoring

Diffstat:
M
tests/sshserver.py
 | 
2
+-
M
tests/test_app.py
 | 
5
+++--
M
webssh/handler.py
 | 
6
+++---
M
webssh/utils.py
 | 
4
++--

4 files changed, 9 insertions(+), 8 deletions(-)
diff --git a/tests/sshserver.py b/tests/sshserver.py
@@ -65,7 +65,7 @@ class Server(paramiko.ServerInterface):
         self.key_verified = False
 
     def get_cmd2enc(self, encodings):
-        while len(encodings) < 2:
+        while len(encodings) < len(self.commands):
             encodings.append(random.choice(self.encodings))
         return dict(zip(self.commands, encodings[0:2]))
 
diff --git a/tests/test_app.py b/tests/test_app.py
@@ -7,7 +7,7 @@ import tornado.gen
 from tornado.testing import AsyncHTTPTestCase
 from tornado.httpclient import HTTPError
 from tornado.options import options
-from tests.sshserver import run_ssh_server, banner
+from tests.sshserver import run_ssh_server, banner, Server
 from tests.utils import encode_multipart_formdata, read_file, make_tests_data_path  # noqa
 from webssh import handler
 from webssh.main import make_app, make_handlers
@@ -25,6 +25,7 @@ except ImportError:
 
 handler.DELAY = 0.1
 swallow_http_errors = handler.swallow_http_errors
+server_encodings = {e.strip() for e in Server.encodings}
 
 
 class TestAppBase(AsyncHTTPTestCase):
@@ -775,7 +776,7 @@ class TestAppWithBadEncoding(OtherTestBase):
         response = yield self.async_post('/', self.body)
         dic = json.loads(to_str(response.body))
         self.assert_status_none(dic)
-        self.assertIn(dic['encoding'], ['UTF-8', 'GBK'])
+        self.assertIn(dic['encoding'], server_encodings)
 
 
 class TestAppWithUnknownEncoding(OtherTestBase):
diff --git a/webssh/handler.py b/webssh/handler.py
@@ -392,7 +392,7 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
 
     def parse_encoding(self, data):
         try:
-            encoding = to_str(data, 'ascii')
+            encoding = to_str(data.strip(), 'ascii')
         except UnicodeDecodeError:
             return
 
@@ -407,8 +407,8 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
 
         for command in commands:
             _, stdout, _ = ssh.exec_command(command, get_pty=True)
-            data = stdout.read().strip()
-            logging.debug('encoding: {}'.format(data))
+            data = stdout.read()
+            logging.debug('{!r} => {!r}'.format(command, data))
             result = self.parse_encoding(data)
             if result:
                 return result
diff --git a/webssh/utils.py b/webssh/utils.py
@@ -54,9 +54,9 @@ def is_valid_port(port):
     return 0 < port < 65536
 
 
-def is_valid_encoding(encoding, ustr=u'test'):
+def is_valid_encoding(encoding):
     try:
-        ustr.encode(encoding)
+        u'test'.encode(encoding)
     except LookupError:
         return False
     return True