完成session共享测试

coco/app.py

@@ -2,6 +2,7 @@ import os
 import time
 import threading
 from queue import Queue
+import logging
 
 from .config import Config
 from .sshd import SSHServer
@@ -11,6 +12,7 @@ from .logging import create_logger
 __version__ = '0.4.0'
 
 BASE_DIR = os.path.dirname(os.path.dirname(__file__))
+logger = logging.getLogger(__file__)
 
 
 class Coco:
@@ -95,20 +97,23 @@ class Coco:
         pass
 
     def shutdown(self):
-        print("Grace shutdown the server")
+        for client in self.clients:
+            self.remove_client(client)
+        time.sleep(1)
         self.sshd.shutdown()
+        logger.info("Grace shutdown the server")
 
     def add_client(self, client):
         with self.lock:
             self.clients.append(client)
-            print("%s add client, now %d s" % (self.name, len(self.clients)))
+            logger.info("New client %s join, total %d now" % (client, len(self.clients)))
 
     def remove_client(self, client):
         with self.lock:
+            try:
                 self.clients.remove(client)
-            print("%s remove client, now %d s" % (self.name, len(self.clients)))
+                logger.info("Client %s leave, total %d now" % (client, len(self.clients)))
 
-        try:
                 client.send("Closed by server")
                 client.close()
             except:

coco/exception.py

+# coding: utf-8
+
+
+class PermissionFailed(Exception):
+    pass

coco/forward.py

@@ -3,21 +3,37 @@
 import socket
 
 import paramiko
+import time
 
 from .session import Session
+from .models import Server
+from .exception import PermissionFailed
 
 
 class ProxyServer:
-    def __init__(self, app, request, client):
+    def __init__(self, app, client, request):
         self.app = app
         self.request = request
         self.client = client
         self.server = None
 
     def proxy(self, asset, system_user):
+        try:
             self.server = self.get_server_conn(asset, system_user)
+        except PermissionFailed:
+            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)
 
     def validate_permission(self, asset, system_user):
         """
@@ -25,26 +41,37 @@ class ProxyServer:
         system user
         :return: True or False
         """
-        pass
+        return True
 
-    def get_system_user_info(self, system_user):
+    def get_system_user_auth(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()
+        if not self.validate_permission(asset, system_user):
+            raise PermissionFailed
+
+        self.get_system_user_auth(system_user)
+        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, paramiko.ssh_exception.NoValidConnectionsError) as e:
-            pass
-
-        except socket.error:
-            pass
+            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:
+            self.client.send("Authentication failed: %s" % e)
+            return
 
+        except socket.error as e:
+            self.client.send("Connection server error: %s" % e)
             return
 
+        term = self.request.meta.get('term', 'xterm')
+        width = self.request.meta.get('width', 80)
+        height = self.request.meta.get('height', 24)
+        chan = ssh.invoke_shell(term, width=width, height=height)
+        return Server(chan, asset, system_user)
+

coco/interactive.py

@@ -79,20 +79,25 @@ class InteractiveServer:
                 input_data.append(data)
 
     def dispatch(self, opt):
-        if opt in ['q', 'Q']:
+        if opt in ['q', 'Q', '0']:
             self.app.remove_client(self.client)
-            return
+        elif opt in ['h', 'H', '9']:
+            self.display_banner()
+        elif opt in ['p', 'P', '3']:
+            self.display_assets()
+        elif opt in ['g', 'G', '5']:
+            self.display_asset_groups()
         else:
-            self.client.send("hello")
-
-        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)
+            self.search_and_proxy(opt)
 
     def search_assets(self, opt, from_result=False):
         pass
 
     def display_assets(self):
+        """
+        Display user all assets
+        :return:
+        """
         pass
 
     def display_asset_groups(self):
@@ -118,7 +123,9 @@ class InteractiveServer:
         pass
 
     def search_and_proxy(self, opt, from_result=False):
-        pass
+        asset = Asset(id=1, hostname="testserver", ip="123.57.183.135", port=8022)
+        system_user = SystemUser(id=2, username="web", password="redhat123", name="web")
+        self.connect(asset, system_user)
 
     def connect(self, asset, system_user):
         forwarder = ProxyServer(self.app, self.client, self.request)

coco/logging.py

@@ -26,8 +26,11 @@ def create_logger(app):
     log_path = os.path.join(log_dir, 'coco.log')
     logger = logging.getLogger()
 
+    # main_formatter = logging.Formatter(
+    #     fmt='%(asctime)s [%(module)s %(levelname)s] %(message)s',
+    #     datefmt='%Y-%m-%d %H:%M:%S')
     main_formatter = logging.Formatter(
-        fmt='%(asctime)s [%(module)s %(levelname)s] %(message)s',
+        fmt='%(asctime)s [%(levelname)s] %(message)s',
         datefmt='%Y-%m-%d %H:%M:%S')
     console_handler = StreamHandler()
     file_handler = TimedRotatingFileHandler(

coco/models.py

@@ -99,7 +99,7 @@ class Server:
         self.system_user = system_user
 
     def fileno(self):
-        return self.chan.fileno
+        return self.chan.fileno()
 
     def send(self, b):
         return self.chan.send(b)
@@ -107,6 +107,10 @@ class Server:
     def recv(self, size):
         return self.chan.recv(size)
 
+    def close(self):
+        self.chan.close()
+        return self.chan.transport.close()
+
     def __getattr__(self, item):
         return getattr(self.chan, item)
 

coco/session.py

@@ -30,51 +30,63 @@ class Session:
         :return:
         """
         logger.info("Session % add watcher %s" % (self, watcher))
+        watcher.send("Welcome to join session %s\r\n" % self.id)
         self.watchers.append(watcher)
 
+    def remove_watcher(self, watcher):
+        logger.info("Session %s remove watcher %s" % (self, watcher))
+        watcher.send("Leave session %s at %s" % (self.id, datetime.datetime.now()))
+        self.watchers.remove(watcher)
+
     def add_sharer(self, sharer):
         """
         Add a sharer, it can read and write to server
         :param sharer:  A client socket
         :return:
         """
-        logger.info("Session % add share %s" % (self, sharer))
+        logger.info("Session %s add share %s" % (self.id, sharer))
+        sharer.send("Welcome to join session %s\r\n" % self.id)
         self.sharers.append(sharer)
 
+    def remove_sharer(self, sharer):
+        logger.info("Session %s remove sharer %s" % (self.id, sharer))
+        sharer.send("Leave session %s at %s" % (self.id, datetime.datetime.now()))
+        self.sharers.remove(sharer)
+
     def bridge(self):
         """
         Bridge clients with server
         :return:
 
         """
-
+        logger.info("Start bridge session %s" % self.id)
         while self.running:
-            try:
-                r, w, x = select.select([self.client + self.server]
-                                 + self.sharers, [], [])
-
+            r, w, x = select.select([self.client, self.server] +
+                                    self.watchers + self.sharers, [], [])
             for sock in r:
-                    if sock == self.server:
                 data = sock.recv(BUF_SIZE)
+                print(data.decode('utf-8'))
+                if sock == self.server:
                     if len(data) == 0:
                         self.close()
+                        break
                     for watcher in [self.client] + self.watchers + self.sharers:
                         watcher.send(data)
                 elif sock == self.client:
-                        data = sock.recv(BUF_SIZE)
                     if len(data) == 0:
                         for watcher in self.watchers + self.sharers:
-                                watcher.send("%s close the session" % self.client)
+                            watcher.send("Client %s close the session" % self.client)
                         self.close()
+                        break
                     self.server.send(data)
                 elif sock in self.sharers:
-                        data = sock.recv(BUF_SIZE)
                     if len(data) == 0:
-                            sock.send("Leave session %s" % self.id)
+                        logger.info("Sharer %s leave session %s" % (sock, self.id))
+                        self.remove_sharer(sock)
                     self.server.send(data)
-
-            except Exception as e:
-                pass
+                elif sock in self.watchers:
+                    if len(data) == 0:
+                        logger.info("Watcher %s leave session %s" % (sock, self.id))
 
     def set_size(self, width, height):
         self.server.resize_pty(width=width, height=height)
@@ -91,10 +103,13 @@ class Session:
         pass
 
     def close(self):
-        pass
+        self.running = False
+        self.server.close()
+        return
 
     def __str__(self):
         return self.id
+    __repr__ = __str__
 
 
 

coco/sshd.py

@@ -87,6 +87,8 @@ class SSHServer:
         except paramiko.SSHException:
             logger.warning("SSH negotiation failed.")
             sys.exit(1)
+        except EOFError:
+            logger.warning("EOF Error")
 
         chan = transport.accept(10)
         if chan is None:

manage.py

@@ -2,9 +2,15 @@
 #
 
 import os
+import sys
 
 from coco import Coco
-import conf
+
+try:
+    import conf
+except ImportError:
+    print("Please prepare config file `cp conf_example.py conf.py`")
+    sys.exit(1)
 
 try:
     os.mkdir("logs")