添加forwarder and some model

coco/forward.py

+#coding: utf-8
+import socket
+
+import paramiko
+
+from .session import Session
+
+
+class ProxyServer:
+    def __init__(self, app, request, client):
+        self.app = app
+        self.request = request
+        self.client = client
+        self.server = None
+
+    def proxy(self, asset, system_user):
+        self.server = self.get_server_conn(asset, system_user)
+        session = Session(self.client, self.server)
+        session.bridge()
+
+    def validate_permission(self, asset, system_user):
+        """
+        Validate use is have the permission to connect this asset using that
+        system user
+        :return: True or False
+        """
+        pass
+
+    def get_system_user_info(self, system_user):
+        """
+        Get the system user auth ..., using this to connect asset
+        :return: system user have full info
+        """
+        pass
+
+    def get_server_conn(self, asset, system_user):
+
+        self.ssh = ssh = paramiko.SSHClient()
+        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        try:
+            ssh.connect(asset.ip, port=asset.port, username=system_user.username,
+                    password=system_user.password, pkey=system_user.private_key)
+        except (paramiko.AuthenticationException,) as e:
+            pass
+
+        except socket.error:
+            pass
+
+        return
+

coco/interactive.py

 #!coding: utf-8
 import socket
 
+import select
+
 from . import char
+from .utils import TtyIOParser, wrap_with_line_feed as wr, \
+    wrap_with_primary as primary, wrap_with_warning as warning
+from .forward import ProxyServer
+from .models import Asset, SystemUser
 
 
 class InteractiveServer:
 
-    def __init__(self, app, request, chan):
+    def __init__(self, app, request, client):
         self.app = app
         self.request = request
-        self.client = chan
+        self.client = client
 
     def display_banner(self):
         self.client.send(char.CLEAR_CHAR)
@@ -27,17 +33,98 @@ class InteractiveServer:
         0) 输入 \033[32mQ/q\033[0m 退出.\r\n""" % self.request.user
         self.client.send(banner)
 
-    def get_input(self, prompt='Opt> '):
+    def get_choice(self, prompt='Opt> '):
+        """实现了一个ssh input, 提示用户输入, 获取并返回
+
+        :return user input string
+        """
+        # Todo: 实现自动hostname或IP补全
+        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))
+        while True:
+            r, w, x = select.select([self.client], [], [])
+            if self.client in r:
+                data = self.client.recv(10)
+                # Client input backspace
+                if data in char.BACKSPACE_CHAR:
+                    # If input words less than 0, should send 'BELL'
+                    if len(input_data) > 0:
+                        data = char.BACKSPACE_CHAR[data]
+                        input_data.pop()
+                    else:
+                        data = char.BELL_CHAR
+                    self.client.send(data)
+                    continue
+
+                # Todo: Move x1b to char
+                if data.startswith(b'\x1b') or data in char.UNSUPPORTED_CHAR:
+                    self.client.send('')
+                    continue
+
+                # handle shell expect
+                multi_char_with_enter = False
+                if len(data) > 1 and data[-1] in char.ENTER_CHAR:
+                    self.client.send(data)
+                    input_data.append(data[:-1])
+                    multi_char_with_enter = True
+
+                # 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))
+                    option = parser.parse_input(b''.join(input_data))
+                    return option.strip()
+                else:
+                    self.client.send(data)
+                    input_data.append(data)
+
+    def dispatch(self, opt):
+        asset = Asset(id=1, hostname="123.57.183.135", ip="123.57.183.135", port=8022)
+        system_user = SystemUser(id=2, username="web", password="redhat123", name="web")
+        self.connect(asset, system_user)
+
+    def search_assets(self, opt, from_result=False):
+        pass
+
+    def display_assets(self):
+        pass
+
+    def display_asset_groups(self):
         pass
 
-    def dispatch(self):
+    def display_group_assets(self):
         pass
 
-    def active(self):
+    def display_search_result(self):
+        pass
+
+    def search_and_display(self, opt):
+        self.search_assets(opt=opt)
+        self.display_search_result()
+
+    def get_user_asset_groups(self):
+        pass
+
+    def get_user_assets(self):
+        pass
+
+    def choose_system_user(self, system_users):
+        pass
+
+    def search_and_proxy(self, opt, from_result=False):
+        pass
+
+    def connect(self, asset, system_user):
+        forwarder = ProxyServer(self.app, self.client, self.request)
+        forwarder.proxy(asset, system_user)
+
+    def activate(self):
         self.display_banner()
         while True:
             try:
-                self.dispatch()
+                opt = self.get_choice()
+                self.dispatch(opt)
             except socket.error:
                 break
         self.close()

coco/models.py

+import json
+import threading
+
+import select
+
+
+BUF_SIZE = 4096
+
+
+class Decoder:
+    def __init__(self, **kwargs):
+        for attr, val in kwargs.items():
+            setattr(self, attr, val)
+
+    @classmethod
+    def from_json(cls, json_str):
+        json_dict = json.loads(json_str)
+        return cls(**json_dict)
+
+    @classmethod
+    def from_multi_json(cls, json_list):
+        json_dict_list = json.loads(json_list)
+        return [cls(**json_dict) for json_dict in json_dict_list]
+
+
+class User(Decoder):
+    id = ""
+    username = ""
+    name = ""
+
+    def __str__(self):
+        return self.name
+    __repr__ = __str__
+
+
+class Asset(Decoder):
+    id = ""
+    hostname = ""
+    ip = ""
+    port = 22
+
+    def __str__(self):
+        return self.hostname
+    __repr__ = __str__
+
+
+class SystemUser(Decoder):
+    id = ""
+    name = ""
+    username = ""
+    password = ""
+    private_key = ""
+
+    def __str__(self):
+        return self.name
+    __repr__ = __str__
+
+
+class Client:
+    """
+    Client is the request client. Nothing more to say
+
+    ```
+    client = Client(chan, addr, user)
+    ```
+
+    """
+
+    def __init__(self, chan, addr, user):
+        self.chan = chan
+        self.addr = addr
+        self.user = user
+
+    def fileno(self):
+        return self.chan.fileno()
+
+    def send(self, b):
+        return self.chan.send(b)
+
+    def recv(self, size):
+        return self.chan.recv(size)
+
+    def __getattr__(self, item):
+        return getattr(self.chan, item)
+
+    def __str__(self):
+        return "<%s from %s:%s>" % (self.user, self.addr[0], self.addr[1])
+
+
+class Server:
+    """
+    Server object like client, a wrapper object, a connection to the asset,
+    Because we don't want to using python dynamic feature, such asset
+    have the chan and system_user attr.
+
+    """
+    # Todo: Server name is not very proper
+    def __init__(self, chan, asset, system_user):
+        self.chan = chan
+        self.asset = asset
+        self.system_user = system_user
+
+    def fileno(self):
+        return self.chan.fileno
+
+    def send(self, b):
+        return self.chan.send(b)
+
+    def recv(self, size):
+        return self.chan.recv(size)
+
+    def __getattr__(self, item):
+        return getattr(self.chan, item)
+
+    def __str__(self):
+        return "<%s@%s:%s>" % \
+               (self.system_user.username, self.asset.hostname, self.asset.port)
+
+
+class WSProxy:
+    """
+    WSProxy is websocket proxy channel object.
+
+    Because tornado or flask websocket base event, if we want reuse func
+    with sshd, we need change it to socket, so we implement a proxy.
+
+    we should use socket pair implement it. usage:
+
+    ```
+
+    child, parent = socket.socketpair()
+
+    # self must have write_message method, write message to ws
+    proxy = WSProxy(self, child)
+    client = Client(parent, user)
+
+    ```
+    """
+
+    def __init__(self, ws, child):
+        """
+        :param ws: websocket instance or handler, have write_message method
+        :param child: sock child pair
+        """
+        self.ws = ws
+        self.child = child
+        self.stop_event = threading.Event()
+
+        self.auto_forward()
+
+    def send(self, b):
+        """
+        If ws use proxy send data, then send the data to child sock, then
+        the parent sock recv
+
+        :param b: data
+        :return:
+        """
+        self.child.send(b)
+
+    def forward(self):
+        while not self.stop_event.is_set():
+            r, w, e = select.select([self.child], [], [])
+            if self.child in r:
+                data = self.child.recv(BUF_SIZE)
+                if len(data) == 0:
+                    self.close()
+                self.ws.write_message(data)
+
+    def auto_forward(self):
+        thread = threading.Thread(target=self.forward, args=())
+        thread.daemon = True
+        thread.start()
+
+    def close(self):
+        self.stop_event.set()
+        self.ws.close()
+        self.child.close()

coco/sshd.py

 #! coding: utf-8
-
+import datetime
 import os
 import logging
 import socket
@@ -10,6 +10,7 @@ import sys
 from .utils import ssh_key_gen
 from .interface import SSHInterface
 from .interactive import InteractiveServer
+from .models import Client
 
 logger = logging.getLogger(__file__)
 BACKLOG = 5
@@ -23,6 +24,7 @@ class Request:
         self.addr = addr
         self.user = None
         self.change_size_event = threading.Event()
+        self.date_start = datetime.datetime.now()
 
 
 class SSHServer:
@@ -57,30 +59,30 @@ class SSHServer:
         max_conn_num = self.app.config['MAX_CONNECTIONS']
         while not self.stop_event.is_set():
             try:
-                client, addr = self.sock.accept()
+                sock, addr = self.sock.accept()
                 logger.info("Get ssh request from %s: %s" % (addr[0], addr[1]))
                 if len(self.app.connections) >= max_conn_num:
-                    client.close()
+                    sock.close()
                     logger.warning("Arrive max connection number %s, "
                                    "reject new request %s:%s" %
                                    (max_conn_num, addr[0], addr[1]))
                 else:
-                    self.app.connections.append((client, addr))
-                thread = threading.Thread(target=self.handle, args=(client, addr))
+                    self.app.connections.append((sock, addr))
+                thread = threading.Thread(target=self.handle, args=(sock, addr))
                 thread.daemon = True
                 thread.start()
             except Exception as e:
                 logger.error("SSH server error: %s" % e)
 
-    def handle(self, client, addr):
-        transport = paramiko.Transport(client, gss_kex=False)
+    def handle(self, sock, addr):
+        transport = paramiko.Transport(sock, gss_kex=False)
         try:
             transport.load_server_moduli()
         except IOError:
             logger.warning("Failed load moduli -- gex will be unsupported")
 
         transport.add_server_key(self.host_key)
-        request = Request(client, addr)
+        request = Request(sock, addr)
         server = SSHInterface(self.app, request)
         try:
             transport.start_server(server=server)
@@ -98,17 +100,18 @@ class SSHServer:
             logger.warning("Client not request a valid request")
             sys.exit(2)
 
-        self.dispatch(request, chan)
+        client = Client(chan, addr, request.user)
+        self.dispatch(request, client)
 
-    def dispatch(self, request, channel):
+    def dispatch(self, request, client):
         if request.type == 'pty':
-            InteractiveServer(self.app, request, channel).active()
+            InteractiveServer(self.app, request, client).activate()
         elif request.type == 'exec':
             pass
         elif request.type == 'subsystem':
             pass
         else:
-            channel.send("Not support request type: %s" % request.type)
+            client.send("Not support request type: %s" % request.type)
 
     def shutdown(self):
         self.stop_event.set()

coco/utils.py

@@ -166,3 +166,59 @@ class TtyIOParser(object):
         self.screen.reset()
         command = self.clean_ps1_etc(command)
         return command
+
+
+def wrap_with_line_feed(s, before=0, after=1):
+    return '\r\n' * before + s + '\r\n' * after
+
+
+def wrap_with_color(text, color='white', background=None,
+                    bolder=False, underline=False):
+    bolder_ = '1'
+    underline_ = '4'
+    color_map = {
+        'black': '30',
+        'red': '31',
+        'green': '32',
+        'brown': '33',
+        'blue': '34',
+        'purple': '35',
+        'cyan': '36',
+        'white': '37',
+    }
+    background_map = {
+        'black': '40',
+        'red': '41',
+        'green': '42',
+        'brown': '43',
+        'blue': '44',
+        'purple': '45',
+        'cyan': '46',
+        'white': '47',
+    }
+
+    wrap_with = []
+    if bolder:
+        wrap_with.append(bolder_)
+    if underline:
+        wrap_with.append(underline_)
+    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'
+
+
+def wrap_with_warning(text, bolder=False):
+    return wrap_with_color(text, color='red', bolder=bolder)
+
+
+def wrap_with_info(text, bolder=False):
+    return wrap_with_color(text, color='brown', bolder=bolder)
+
+
+def wrap_with_primary(text, bolder=False):
+    return wrap_with_color(text, color='green', bolder=bolder)
+
+
+def wrap_with_title(text):
+    return wrap_with_color(text, color='black', background='green')