diff --git a/tests/sshserver.py b/tests/sshserver.py index 9961fdb..490df9e 100644 --- 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 index ff8e399..a88172a 100644 --- 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 index dd5050f..a852acc 100644 --- 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 index 3bdfda4..845ca56 100644 --- 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