Auto detect system default encoding

tests/sshserver.py

@@ -22,6 +22,7 @@ from binascii import hexlify
 import socket
 # import sys
 import threading
+import random
 # import traceback
 
 import paramiko
@@ -36,6 +37,8 @@ host_key = paramiko.RSAKey(filename='tests/test_rsa.key')
 
 print('Read key: ' + u(hexlify(host_key.get_fingerprint())))
 
+banner = u'\r\n\u6b22\u8fce\r\n'
+
 
 class Server(paramiko.ServerInterface):
     # 'data' is the output of base64.b64encode(key)
@@ -46,8 +49,13 @@ class Server (paramiko.ServerInterface):
             b'UWT10hcuO4Ks8=')
     good_pub_key = paramiko.RSAKey(data=decodebytes(data))
 
+    langs = ['en_US.UTF-8', 'zh_CN.GBK']
+
     def __init__(self):
-        self.event = threading.Event()
+        self.shell_event = threading.Event()
+        self.exec_event = threading.Event()
+        self.lang = random.choice(self.langs)
+        self.encoding = self.lang.split('.')[-1]
 
     def check_channel_request(self, kind, chanid):
         if kind == 'session':
@@ -68,8 +76,19 @@ class Server (paramiko.ServerInterface):
     def get_allowed_auths(self, username):
         return 'password,publickey'
 
+    def check_channel_exec_request(self, channel, command):
+        if command != b'locale':
+            ret = False
+        else:
+            ret = True
+            result = 'LANG={lang}\nLANGUAGE=\nLC_CTYPE="{lang}"\n'.format(lang=self.lang)  # noqa
+            channel.send(result)
+            channel.shutdown(1)
+        self.exec_event.set()
+        return ret
+
     def check_channel_shell_request(self, channel):
-        self.event.set()
+        self.shell_event.set()
         return True
 
     def check_channel_pty_request(self, channel, term, width, height,
@@ -112,12 +131,19 @@ def run_ssh_server(port=2200, running=True):
         username = t.get_username()
         print('{} Authenticated!'.format(username))
 
-        server.event.wait(10)
-        if not server.event.is_set():
+        server.shell_event.wait(2)
+        if not server.shell_event.is_set():
             print('*** Client never asked for a shell.')
             continue
 
-        chan.send('\r\n\r\nWelcome!\r\n\r\n')
+        server.exec_event.wait(2)
+        if not server.exec_event.is_set():
+            print('*** Client never asked for a command.')
+            continue
+
+        # chan.send('\r\n\r\nWelcome!\r\n\r\n')
+        print(server.encoding)
+        chan.send(banner.encode(server.encoding))
         if username == 'bar':
             msg = chan.recv(1024)
             chan.send(msg)

tests/test_app.py

@@ -9,7 +9,7 @@ from tornado.testing import AsyncHTTPTestCase
 from tornado.options import options
 from webssh.main import make_app, make_handlers
 from webssh.settings import get_app_settings
-from tests.sshserver import run_ssh_server
+from tests.sshserver import run_ssh_server, banner
 
 
 handler.DELAY = 0.1
@@ -79,8 +79,10 @@ class TestApp(AsyncHTTPTestCase):
         response = self.fetch('/')
         self.assertEqual(response.code, 200)
         response = self.fetch('/', method="POST", body=self.body)
-        worker_id = json.loads(response.body.decode('utf-8'))['id']
-        self.assertIsNotNone(worker_id)
+        data = json.loads(response.body.decode('utf-8'))
+        self.assertIsNone(data['status'])
+        self.assertIsNotNone(data['id'])
+        self.assertIsNotNone(data['encoding'])
 
     @tornado.testing.gen_test
     def test_app_with_correct_credentials_timeout(self):
@@ -90,11 +92,13 @@ class TestApp(AsyncHTTPTestCase):
         self.assertEqual(response.code, 200)
 
         response = yield client.fetch(url, method="POST", body=self.body)
-        worker_id = json.loads(response.body.decode('utf-8'))['id']
-        self.assertIsNotNone(worker_id)
+        data = json.loads(response.body.decode('utf-8'))
+        self.assertIsNone(data['status'])
+        self.assertIsNotNone(data['id'])
+        self.assertIsNotNone(data['encoding'])
 
         url = url.replace('http', 'ws')
-        ws_url = url + 'ws?id=' + worker_id
+        ws_url = url + 'ws?id=' + data['id']
         yield tornado.gen.sleep(handler.DELAY + 0.1)
         ws = yield tornado.websocket.websocket_connect(ws_url)
         msg = yield ws.read_message()
@@ -109,14 +113,16 @@ class TestApp(AsyncHTTPTestCase):
         self.assertEqual(response.code, 200)
 
         response = yield client.fetch(url, method="POST", body=self.body)
-        worker_id = json.loads(response.body.decode('utf-8'))['id']
-        self.assertIsNotNone(worker_id)
+        data = json.loads(response.body.decode('utf-8'))
+        self.assertIsNone(data['status'])
+        self.assertIsNotNone(data['id'])
+        self.assertIsNotNone(data['encoding'])
 
         url = url.replace('http', 'ws')
-        ws_url = url + 'ws?id=' + worker_id
+        ws_url = url + 'ws?id=' + data['id']
         ws = yield tornado.websocket.websocket_connect(ws_url)
         msg = yield ws.read_message()
-        self.assertIn(b'Welcome!', msg)
+        self.assertEqual(msg.decode(data['encoding']), banner)
         ws.close()
 
     @tornado.testing.gen_test
@@ -128,14 +134,16 @@ class TestApp(AsyncHTTPTestCase):
 
         body = self.body.replace('robey', 'bar')
         response = yield client.fetch(url, method="POST", body=body)
-        worker_id = json.loads(response.body.decode('utf-8'))['id']
-        self.assertIsNotNone(worker_id)
+        data = json.loads(response.body.decode('utf-8'))
+        self.assertIsNone(data['status'])
+        self.assertIsNotNone(data['id'])
+        self.assertIsNotNone(data['encoding'])
 
         url = url.replace('http', 'ws')
-        ws_url = url + 'ws?id=' + worker_id
+        ws_url = url + 'ws?id=' + data['id']
         ws = yield tornado.websocket.websocket_connect(ws_url)
         msg = yield ws.read_message()
-        self.assertIn(b'Welcome!', msg)
+        self.assertEqual(msg.decode(data['encoding']), banner)
 
         # messages below will be ignored silently
         yield ws.write_message('hello')

tests/test_handler.py

@@ -3,7 +3,22 @@ import os.path
 import paramiko
 
 from tornado.httputil import HTTPServerRequest
-from webssh.handler import MixinHandler, IndexHandler
+from webssh.handler import MixinHandler, IndexHandler, parse_encoding
+
+
+class TestHandler(unittest.TestCase):
+
+    def test_parse_encoding(self):
+        data = ''
+        self.assertIsNone(parse_encoding(data))
+        data = 'UTF-8'
+        self.assertEqual(parse_encoding(data), 'UTF-8')
+        data = 'en_US.UTF-8'
+        self.assertEqual(parse_encoding(data), 'UTF-8')
+        data = 'LANG=en_US.UTF-8\nLANGUAGE=\nLC_CTYPE="en_US.UTF-8"\n'
+        self.assertEqual(parse_encoding(data), 'UTF-8')
+        data = 'LANGUAGE=\nLC_CTYPE="en_US.UTF-8"\n'
+        self.assertEqual(parse_encoding(data), 'UTF-8')
 
 
 class TestMixinHandler(unittest.TestCase):

webssh/handler.py

@@ -27,6 +27,13 @@ except ImportError:
 DELAY = 3
 
 
+def parse_encoding(data):
+    for line in data.split('\n'):
+        s = line.split('=')[-1]
+        if s:
+            return s.strip('"').split('.')[-1]
+
+
 class MixinHandler(object):
 
     def get_real_client_addr(self):
@@ -122,6 +129,17 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
         return self.get_real_client_addr() or self.request.connection.stream.\
             socket.getpeername()
 
+    def get_default_encoding(self, ssh):
+        try:
+            _, stdout, _ = ssh.exec_command('locale')
+        except paramiko.SSHException:
+            result = None
+        else:
+            data = stdout.read().decode()
+            result = parse_encoding(data)
+
+        return result if result else 'utf-8'
+
     def ssh_connect(self):
         ssh = paramiko.SSHClient()
         ssh._system_host_keys = self.host_keys_settings['system_host_keys']
@@ -146,6 +164,7 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
         chan.setblocking(0)
         worker = Worker(self.loop, ssh, chan, dst_addr)
         worker.src_addr = self.get_client_addr()
+        worker.encoding = self.get_default_encoding(ssh)
         return worker
 
     def ssh_connect_wrapped(self, future):
@@ -164,6 +183,7 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
     def post(self):
         worker_id = None
         status = None
+        encoding = None
 
         future = Future()
         t = threading.Thread(target=self.ssh_connect_wrapped, args=(future,))
@@ -178,8 +198,9 @@ class IndexHandler(MixinHandler, tornado.web.RequestHandler):
             worker_id = worker.id
             workers[worker_id] = worker
             self.loop.call_later(DELAY, recycle_worker, worker)
+            encoding = worker.encoding
 
-        self.write(dict(id=worker_id, status=status))
+        self.write(dict(id=worker_id, status=status, encoding=encoding))
 
 
 class WsockHandler(MixinHandler, tornado.websocket.WebSocketHandler):

webssh/static/js/main.js

@@ -59,12 +59,14 @@ jQuery(function($){
         join = (ws_url[ws_url.length-1] === '/' ? '' : '/'),
         url = ws_url + join + 'ws?id=' + msg.id,
         sock = new window.WebSocket(url),
+        encoding = msg.encoding,
         terminal = document.getElementById('#terminal'),
         term = new window.Terminal({
           cursorBlink: true,
         });
 
     console.log(url);
+    console.log(encoding);
     wssh.sock = sock;
     wssh.term = term;
 
@@ -83,7 +85,7 @@ jQuery(function($){
       var reader = new window.FileReader();
 
       reader.onloadend = function(){
-        var decoder = new window.TextDecoder();
+        var decoder = new window.TextDecoder(encoding);
         var text = decoder.decode(reader.result);
         // console.log(text);
         term.write(text);