[ISSUE #18] Add custom header support for invoke_shell()

.gitignore

@@ -1,6 +1,8 @@
 *.py[co]
 *.egg
 *.egg-info
+*.tmp
 dist
 build
 eggs
+*.DS_Store

examples/flask_server.py

@@ -52,7 +52,7 @@ def index():
 
 if __name__ == '__main__':
     from gevent.pywsgi import WSGIServer
-    from geventwebsocket import WebSocketHandler
+    from geventwebsocket.handler import WebSocketHandler
 
     app.debug = True
     http_server = WSGIServer(('localhost', 5000), app,

setup.py

@@ -3,8 +3,8 @@
 
 setup(
     name='wssh',
-    version='0.1.0',
-    author='Andrea Luzzardi <aluzzardi@gmail.com>',
+    version='0.2.1',
+    author='EricPai <ericpai94@hotmail.com>',
     packages=[
         'wssh'
         ],

wssh/__init__.py

@@ -1,3 +1,3 @@
 from server import WSSHBridge
 
-__version__ = '0.1.0'
+__version__ = '0.2.1'

wssh/client.py

@@ -1,4 +1,5 @@
 import os
+import codecs
 import sys
 import signal
 import errno
@@ -10,68 +11,69 @@
 import tty
 import fcntl
 import struct
-
 import platform
+import ssl
 
 try:
     import simplejson as json
 except ImportError:
     import json
 
 
-class ConnectionError(Exception):
-    pass
-
-
-def _pty_size():
+def _pty_size(utf_in, utf_out):
     rows, cols = 24, 80
     # Can't do much for Windows
     if platform.system() == 'Windows':
         return rows, cols
     fmt = 'HH'
     buffer = struct.pack(fmt, 0, 0)
-    result = fcntl.ioctl(sys.stdout.fileno(), termios.TIOCGWINSZ,
+    result = fcntl.ioctl(utf_out.fileno(), termios.TIOCGWINSZ,
         buffer)
     rows, cols = struct.unpack(fmt, result)
     return rows, cols
 
 
-def _resize(ws):
-    rows, cols = _pty_size()
+def _resize(ws, utf_in, utf_out):
+    rows, cols = _pty_size(utf_in, utf_out)
     ws.send(json.dumps({'resize': {'width': cols, 'height': rows}}))
 
 
-def invoke_shell(endpoint):
-    ssh = websocket.create_connection(endpoint)
-    _resize(ssh)
-    oldtty = termios.tcgetattr(sys.stdin)
+def invoke_shell(endpoint, header = None):
+    UTF8Reader = codecs.getreader('utf8')
+    utf_in = UTF8Reader(sys.stdin)
+    UTF8Writer = codecs.getwriter('utf8')
+    utf_out = UTF8Writer(sys.stdout)
+    ssh = websocket.create_connection(url = endpoint, header = header, sslopt={"cert_reqs": ssl.CERT_NONE})
+    _resize(ssh, utf_in, utf_out)
+    oldtty = termios.tcgetattr(utf_in)
     old_handler = signal.getsignal(signal.SIGWINCH)
 
     def on_term_resize(signum, frame):
-        _resize(ssh)
+        _resize(ssh, utf_in, utf_out)
     signal.signal(signal.SIGWINCH, on_term_resize)
 
     try:
-        tty.setraw(sys.stdin.fileno())
-        tty.setcbreak(sys.stdin.fileno())
+        tty.setraw(utf_in.fileno())
+        tty.setcbreak(utf_in.fileno())
 
-        rows, cols = _pty_size()
+        rows, cols = _pty_size(utf_in, utf_out)
         ssh.send(json.dumps({'resize': {'width': cols, 'height': rows}}))
 
         while True:
             try:
-                r, w, e = select.select([ssh.sock, sys.stdin], [], [])
+                r, w, e = select.select([ssh.sock, utf_in], [], [])
                 if ssh.sock in r:
                     data = ssh.recv()
                     if not data:
                         break
                     message = json.loads(data)
                     if 'error' in message:
-                        raise ConnectionError(message['error'])
-                    sys.stdout.write(message['data'])
-                    sys.stdout.flush()
-                if sys.stdin in r:
-                    x = sys.stdin.read(1)
+                        utf_out.write(message['error'])
+                        break
+                    utf_out.write(message['data'])
+                    utf_out.flush()
+                if utf_in in r:
+                    x = os.read(utf_in.fileno(), 3)
                     if len(x) == 0:
                         break
                     ssh.send(json.dumps({'data': x}))
@@ -83,5 +85,6 @@ def on_term_resize(signum, frame):
     except websocket.WebSocketException:
         raise
     finally:
-        termios.tcsetattr(sys.stdin, termios.TCSADRAIN, oldtty)
+        termios.tcsetattr(utf_in, termios.TCSADRAIN, oldtty)
         signal.signal(signal.SIGWINCH, old_handler)
+        print '\n'

wssh/server.py

@@ -18,6 +18,7 @@
 from paramiko.ssh_exception import SSHException
 
 import socket
+import time
 
 try:
     import simplejson as json
@@ -111,6 +112,8 @@ def _forward_inbound(self, channel):
             while True:
                 data = self._websocket.receive()
                 if not data:
+                    channel.send('\x03\rexit\r')
+                    time.sleep(10)
                     return
                 data = json.loads(str(data))
                 if 'resize' in data:
@@ -127,10 +130,13 @@ def _forward_outbound(self, channel):
         try:
             while True:
                 wait_read(channel.fileno())
-                data = channel.recv(1024)
+                data = channel.recv(10240)
                 if not len(data):
                     return
                 self._websocket.send(json.dumps({'data': data}))
+        except Exception as e:
+            channel.send('\x03\rexit\r')
+            time.sleep(10)
         finally:
             self.close()
 

wssh/static/term.js

@@ -131,7 +131,8 @@ function Terminal(cols, rows, handler) {
   this.rows = rows || Terminal.geometry[1];
 
   if (handler) {
-    this.on('data', handler);
+    this.dataHandler = handler;
+    this.on('data', this.dataHandler);
   }
 
   this.ybase = 0;
@@ -2432,7 +2433,7 @@ Terminal.prototype.reverseIndex = function() {
 
 // ESC c Full Reset (RIS).
 Terminal.prototype.reset = function() {
-  Terminal.call(this, this.cols, this.rows);
+  Terminal.call(this, this.cols, this.rows, this.dataHandler);
   this.refresh(0, this.rows - 1);
 };