Add websocket server

coco/forward.py

@@ -24,16 +24,10 @@ class ProxyServer:
             self.client.send("No permission")
             return
 
-        if len(self.app.sessions) == 1:
-            session = self.app.sessions[0]
-            session.add_sharer(self.client)
-            while True:
-                time.sleep(10)
-        else:
-            session = Session(self.client, self.server)
-            self.app.sessions.append(session)
-            session.bridge()
-            self.app.sessions.remove(session)
+        session = Session(self.client, self.server)
+        self.app.sessions.append(session)
+        session.bridge()
+        self.app.sessions.remove(session)
 
     def validate_permission(self, asset, system_user):
         """

coco/interactive.py

 #!coding: utf-8
 import socket
+import threading
 
 from . import char
 from .utils import TtyIOParser, wrap_with_line_feed as wr, \
@@ -29,7 +30,7 @@ class InteractiveServer:
         8) 输入 \033[32mD/d\033[0m 批量下载文件.(未完成)\r
         9) 输入 \033[32mH/h\033[0m 帮助.\r
         0) 输入 \033[32mQ/q\033[0m 退出.\r\n""" % self.request.user
-        self.client.send(banner)
+        self.client.send(banner.encode('utf-8'))
 
     def get_choice(self, prompt='Opt> '):
         """实现了一个ssh input, 提示用户输入, 获取并返回
@@ -40,7 +41,7 @@ class InteractiveServer:
         input_data = []
         parser = TtyIOParser(self.request.meta.get("width", 80),
                              self.request.meta.get("height", 24))
-        self.client.send(wr(prompt, before=1, after=0))
+        self.client.send(wr(prompt, before=1, after=0).encode('utf-8'))
         while True:
             data = self.client.recv(10)
             if len(data) == 0:
@@ -59,7 +60,7 @@ class InteractiveServer:
 
             # Todo: Move x1b to char
             if data.startswith(b'\x1b') or data in char.UNSUPPORTED_CHAR:
-                self.client.send('')
+                self.client.send(b'')
                 continue
 
             # handle shell expect
@@ -71,7 +72,7 @@ class InteractiveServer:
 
             # If user type ENTER we should get user input
             if data in char.ENTER_CHAR or multi_char_with_enter:
-                self.client.send(wr('', after=2))
+                self.client.send(wr(b'', after=2))
                 option = parser.parse_input(b''.join(input_data))
                 return option.strip()
             else:
@@ -141,5 +142,10 @@ class InteractiveServer:
                 break
         self.close()
 
+    def activate_async(self):
+        thread = threading.Thread(target=self.activate)
+        thread.daemon = True
+        thread.start()
+
     def close(self):
         pass

coco/models.py

 import json
 import threading
+import datetime
 
 
 BUF_SIZE = 4096
@@ -54,6 +55,16 @@ class SystemUser(Decoder):
     __repr__ = __str__
 
 
+class Request:
+    def __init__(self, remote_ip=""):
+        self.type = ""
+        self.meta = {}
+        self.user = None
+        self.remote_ip = remote_ip
+        self.change_size_event = threading.Event()
+        self.date_start = datetime.datetime.now()
+
+
 class Client:
     """
     Client is the request client. Nothing more to say
@@ -148,6 +159,9 @@ class WSProxy:
         self.child = child
         self.stop_event = threading.Event()
 
+        print(self.child)
+        print(self.child.fileno())
+
         self.auto_forward()
 
     def send(self, b):
@@ -158,7 +172,8 @@ class WSProxy:
         :param b: data
         :return:
         """
-        self.child.send(b)
+        if isinstance(b, str):
+            b = b.encode('utf-8')
 
     def forward(self):
         while not self.stop_event.is_set():

coco/sshd.py

 #! coding: utf-8
-import datetime
+
 import os
 import logging
 import socket
@@ -10,23 +10,12 @@ import sys
 from .utils import ssh_key_gen
 from .interface import SSHInterface
 from .interactive import InteractiveServer
-from .models import Client
+from .models import Client, Request
 
 logger = logging.getLogger(__file__)
 BACKLOG = 5
 
 
-class Request:
-    def __init__(self, client, addr):
-        self.type = ""
-        self.meta = {}
-        self.client = client
-        self.addr = addr
-        self.user = None
-        self.change_size_event = threading.Event()
-        self.date_start = datetime.datetime.now()
-
-
 class SSHServer:
     def __init__(self, app=None):
         self.app = app
@@ -38,7 +27,7 @@ class SSHServer:
     def listen(self):
         host = self.app.config["BIND_HOST"]
         port = self.app.config["SSHD_PORT"]
-        print('Starting shh server at %(host)s:%(port)s' %
+        print('Starting ssh server at %(host)s:%(port)s' %
               {"host": host, "port": port})
         self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         self.sock.bind((host, port))
@@ -80,7 +69,7 @@ class SSHServer:
             logger.warning("Failed load moduli -- gex will be unsupported")
 
         transport.add_server_key(self.host_key)
-        request = Request(sock, addr)
+        request = Request()
         server = SSHInterface(self.app, request)
         try:
             transport.start_server(server=server)

coco/utils.py

@@ -169,6 +169,8 @@ class TtyIOParser(object):
 
 
 def wrap_with_line_feed(s, before=0, after=1):
+    if isinstance(s, bytes):
+        return b'\r\n' * before + s + b'\r\n' * after
     return '\r\n' * before + s + '\r\n' * after
 
 
@@ -205,7 +207,11 @@ def wrap_with_color(text, color='white', background=None,
     if background:
         wrap_with.append(background_map.get(background, ''))
     wrap_with.append(color_map.get(color, ''))
-    return '\033[' + ';'.join(wrap_with) + 'm' + text + '\033[0m'
+
+    data = '\033[' + ';'.join(wrap_with) + 'm' + text + '\033[0m'
+    if isinstance(text, bytes):
+        return data.encode('utf-8')
+    return data
 
 
 def wrap_with_warning(text, bolder=False):

coco/ws.py

 # coding: utf-8
+import socket
+import threading
 
+import time
 import tornado.web
 import tornado.websocket
 import tornado.httpclient
 import tornado.ioloop
+import tornado.gen
 
+from .models import User, Request, Client, WSProxy
+from .interactive import InteractiveServer
 
-class MainHandler(tornado.web.RequestHandler):
-    def get(self, *args, **kwargs):
-        self.write("Hello world")
 
+class BaseWehSocketHandler:
+    def prepare(self):
+        self.app = self.settings["app"]
+        child, parent = socket.socketpair()
+        addr = (self.request.remote_ip, 0)
+        self.client = Client(parent, addr, self.current_user)
+        self.proxy = WSProxy(self, child)
+        self.app.clients.append(self.client)
 
-class InteractiveHandler(tornado.websocket.WebSocketHandler):
-    pass
+    def get_current_user(self):
+        return User(id=1, username="guanghongwei", name="广宏伟")
+
+    def check_origin(self, origin):
+        return True
+
+
+class InteractiveWehSocketHandler(BaseWehSocketHandler, tornado.websocket.WebSocketHandler):
+    @tornado.web.authenticated
+    def open(self):
+        request = Request(self.request.remote_ip)
+        self.request.__dict__.update(request.__dict__)
+        InteractiveServer(self.app, self.request,self.client).activate_async()
 
+    def on_message(self, message):
+        print(message)
+        self.proxy.send(message)
 
-class ProxyHandler(tornado.websocket.WebSocketHandler):
+    def on_close(self):
+        self.proxy.close()
+
+
+class ProxyWehSocketHandler(BaseWehSocketHandler):
     pass
 
 
-class MonitorHandler(tornado.websocket.WebSocketHandler):
+class MonitorWehSocketHandler(BaseWehSocketHandler):
     pass
 
 
 class WSServer:
     routers = [
-        (r'^/$', MainHandler),
-        (r'/ws/interactive/', MainHandler),
-        (r'/ws/proxy/(?P<asset_id>[0-9]+)/(?P<system_user_id>[0-9]+)/', MainHandler),
-        (r'/ws/session/(?P<session_id>[0-9]+)/monitor/', MainHandler),
-        (r'/ws/session/(?P<session+id>[0-9]+)/join/', MainHandler),
+        (r'/ws/interactive/', InteractiveWehSocketHandler),
+        (r'/ws/proxy/(?P<asset_id>[0-9]+)/(?P<system_user_id>[0-9]+)/', ProxyWehSocketHandler),
+        (r'/ws/session/(?P<session_id>[0-9]+)/monitor/', MonitorWehSocketHandler),
     ]
 
+    # prepare may be rewrite it
+    settings = {
+        'cookie_secret': '',
+        'app': None,
+        'login_url': '/login'
+    }
+
     def __init__(self, app):
         self.app = app
+        self._prepare()
+
+    def _prepare(self):
+        self.settings['cookie_secret'] = self.app.config['SECRET_KEY']
+        self.settings['app'] = self.app
 
     def run(self):
         host = self.app.config["BIND_HOST"]
         port = self.app.config["WS_PORT"]
         print('Starting websocket server at %(host)s:%(port)s' %
               {"host": host, "port": port})
-        ws = tornado.web.Application(self.routers)
+        ws = tornado.web.Application(self.routers, **self.settings)
         ws.listen(port=port, address=host)
         tornado.ioloop.IOLoop.current().start()