Fix some bug

coco/forward.py

@@ -5,22 +5,26 @@ import threading
 import logging
 
 import paramiko
+import time
 
 from .session import Session
 from .models import Server
 
 
 logger = logging.getLogger(__file__)
+TIMEOUT = 8
 
 
 class ProxyServer:
-    def __init__(self, app, client, request):
+    def __init__(self, app, client):
         self.app = app
-        self.request = request
         self.client = client
+        self.request = client.request
         self.server = None
+        self.connecting = True
 
     def proxy(self, asset, system_user):
+        self.send_connecting_message()
         self.server = self.get_server_conn(asset, system_user)
         if self.server is None:
             return
@@ -46,6 +50,7 @@ class ProxyServer:
         """
 
     def get_server_conn(self, asset, system_user):
+        logger.info("Connect to %s" % asset.hostname)
         if not self.validate_permission(asset, system_user):
             self.client.send(b'No permission')
             return None
@@ -57,14 +62,18 @@ class ProxyServer:
             ssh.connect(asset.ip, port=asset.port,
                         username=system_user.username,
                         password=system_user.password,
-                        pkey=system_user.private_key)
+                        pkey=system_user.private_key,
+                        timeout=TIMEOUT)
         except paramiko.AuthenticationException as e:
-            self.client.send("Authentication failed: {}".format(e).encode("utf-8"))
+            self.client.send("[Errno 66] Authentication failed: {}".format(e).encode("utf-8"))
             return None
 
         except socket.error as e:
-            self.client.send("Connection server error: {}".format(e).encode("utf-8"))
+            self.client.send(" {}".format(e).encode("utf-8"))
             return None
+        finally:
+            self.connecting = False
+            self.client.send(b'\r\n')
 
         term = self.request.meta.get('term', 'xterm')
         width = self.request.meta.get('width', 80)
@@ -85,3 +94,15 @@ class ProxyServer:
         thread.daemon = True
         thread.start()
 
+    def send_connecting_message(self):
+        def func():
+            delay = 0.0
+            self.client.send('Connecting to {} {:.1f}'.format('abc.com', delay).encode('utf-8'))
+            while self.connecting and delay < TIMEOUT:
+                self.client.send('\x08\x08\x08{:.1f}'.format(delay).encode('utf-8'))
+                time.sleep(0.1)
+                delay += 0.1
+        thread = threading.Thread(target=func)
+        thread.daemon = True
+        thread.start()
+

coco/interactive.py

 #!coding: utf-8
+import logging
 import socket
 import threading
 
@@ -9,13 +10,15 @@ from .forward import ProxyServer
 from .models import Asset, SystemUser
 from .session import Session
 
+logger = logging.getLogger(__file__)
+
 
 class InteractiveServer:
 
-    def __init__(self, app, request, client):
+    def __init__(self, app, client):
         self.app = app
-        self.request = request
         self.client = client
+        self.request = client.request
 
     def display_banner(self):
         self.client.send(char.CLEAR_CHAR)
@@ -40,8 +43,7 @@ class InteractiveServer:
         """
         # Todo: 实现自动hostname或IP补全
         input_data = []
-        parser = TtyIOParser(self.request.meta.get("width", 80),
-                             self.request.meta.get("height", 24))
+        parser = TtyIOParser()
         self.client.send(wr(prompt, before=1, after=0))
         while True:
             data = self.client.recv(10)
@@ -128,10 +130,10 @@ class InteractiveServer:
     def search_and_proxy(self, opt, from_result=False):
         asset = Asset(id=1, hostname="testserver", ip="192.168.244.164", port=22)
         system_user = SystemUser(id=2, username="root", password="redhat", name="web")
-        self.connect(asset, system_user)
+        self.proxy(asset, system_user)
 
-    def connect(self, asset, system_user):
-        forwarder = ProxyServer(self.app, self.client, self.request)
+    def proxy(self, asset, system_user):
+        forwarder = ProxyServer(self.app, self.client)
         forwarder.proxy(asset, system_user)
 
     def replay_session(self, session_id):
@@ -145,7 +147,8 @@ class InteractiveServer:
             try:
                 opt = self.get_choice()
                 self.dispatch(opt)
-            except socket.error:
+            except socket.error as e:
+                logger.error("Cocket error %s" % e)
                 break
         self.close()
 

coco/interface.py

@@ -100,7 +100,7 @@ class SSHInterface(paramiko.ServerInterface):
         return paramiko.OPEN_SUCCEEDED
 
     def check_channel_shell_request(self, channel):
-        logger.info("Check channel shell request: %s" % channel)
+        logger.info("Check channel shell request: %s" % channel.get_id())
         self.event.set()
         return True
 

coco/models.py

@@ -58,11 +58,12 @@ class SystemUser(Decoder):
 
 
 class Request:
-    def __init__(self, remote_ip=""):
+    def __init__(self, addr):
         self.type = ""
         self.meta = {}
         self.user = None
-        self.remote_ip = remote_ip
+        self.addr = addr
+        self.remote_ip = self.addr[0]
         self.change_size_event = threading.Event()
         self.date_start = datetime.datetime.now()
 
@@ -74,13 +75,13 @@ class Client:
     ```
     client = Client(chan, addr, user)
     ```
-
     """
 
-    def __init__(self, chan, addr, user):
+    def __init__(self, chan, request):
         self.chan = chan
-        self.addr = addr
-        self.user = user
+        self.request = request
+        self.user = request.user
+        self.addr = request.addr
 
     def fileno(self):
         return self.chan.fileno()
@@ -119,7 +120,6 @@ class Server:
         self.output = ''
         self._in_input_state = True
         self._input_initial = False
-        self._in_auto_complete_state = False
         self._in_vim_state = False
 
     def fileno(self):
@@ -140,11 +140,6 @@ class Server:
                 del self.output_data[:]
             self._in_input_state = True
 
-        # if b == '\t':
-        #     self._in_auto_complete_state = True
-        # else:
-        #     self._in_auto_complete_state = False
-
         print("Send: %s" % b)
         return self.chan.send(b)
 

coco/session.py

@@ -37,7 +37,7 @@ class Session:
         :param silent: If true not send welcome message
         :return:
         """
-        logger.info("Session % add watcher %s" % (self, watcher))
+        logger.info("Session %s add watcher %s" % (self, watcher))
         if not silent:
             watcher.send("Welcome to watch session {}\r\n".format(self.id).encode("utf-8"))
         self.sel.register(watcher, selectors.EVENT_READ)
@@ -165,8 +165,7 @@ class Session:
 
     def close(self):
         self.running = False
-        for chan in [self.client, self.server] + self.watchers + self.sharers:
-            chan.close()
+        self.server.close()
 
     def __str__(self):
         return self.id

coco/sshd.py

@@ -22,18 +22,9 @@ class SSHServer:
         self.stop_event = threading.Event()
         self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
         self.host_key_path = os.path.join(self.app.root_path, 'keys', 'host_rsa_key')
-        self.host_key = self.get_host_key()
 
-    def listen(self):
-        host = self.app.config["BIND_HOST"]
-        port = self.app.config["SSHD_PORT"]
-        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))
-        self.sock.listen(BACKLOG)
-
-    def get_host_key(self):
+    @property
+    def host_key(self):
         if not os.path.isfile(self.host_key_path):
             self.gen_host_key()
         return paramiko.RSAKey(filename=self.host_key_path)
@@ -44,7 +35,13 @@ class SSHServer:
             f.write(ssh_key)
 
     def run(self):
-        self.listen()
+        host = self.app.config["BIND_HOST"]
+        port = self.app.config["SSHD_PORT"]
+        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))
+        self.sock.listen(BACKLOG)
         while not self.stop_event.is_set():
             try:
                 sock, addr = self.sock.accept()
@@ -63,7 +60,7 @@ class SSHServer:
             logger.warning("Failed load moduli -- gex will be unsupported")
 
         transport.add_server_key(self.host_key)
-        request = Request()
+        request = Request(addr)
         server = SSHInterface(self.app, request)
         try:
             transport.start_server(server=server)
@@ -72,6 +69,7 @@ class SSHServer:
             sys.exit(1)
         except EOFError:
             logger.warning("EOF Error")
+            sys.exit(1)
 
         chan = transport.accept(10)
         if chan is None:
@@ -83,19 +81,20 @@ class SSHServer:
             logger.warning("Client not request a valid request")
             sys.exit(2)
 
-        client = Client(chan, addr, request.user)
+        client = Client(chan, request)
         self.app.add_client(client)
-        self.dispatch(request, client)
+        self.dispatch(client)
 
-    def dispatch(self, request, client):
-        if request.type == 'pty':
-            InteractiveServer(self.app, request, client).activate()
-        elif request.type == 'exec':
+    def dispatch(self, client):
+        request_type = client.request.type
+        if request_type == 'pty':
+            InteractiveServer(self.app, client).activate()
+        elif request_type == 'exec':
             pass
-        elif request.type == 'subsystem':
+        elif request_type == 'subsystem':
             pass
         else:
-            client.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/ws.py

@@ -20,8 +20,10 @@ 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)
+        request = Request((self.request.remote_ip, 0))
+        request.user = self.current_user
+        self.request.__dict__.update(request.__dict__)
+        self.client = Client(parent, self.request)
         self.proxy = WSProxy(self, child)
         self.app.clients.append(self.client)
 
@@ -35,9 +37,7 @@ class BaseWehSocketHandler:
 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()
+        InteractiveServer(self.app, self.client).activate_async()
 
     def on_message(self, message):
         try: