Add command parser

coco/exception.py

@@ -3,3 +3,4 @@
 
 class PermissionFailed(Exception):
     pass
+

coco/forward.py

@@ -8,7 +8,6 @@ import paramiko
 
 from .session import Session
 from .models import Server
-from .exception import PermissionFailed
 
 
 logger = logging.getLogger(__file__)
@@ -22,12 +21,9 @@ class ProxyServer:
         self.server = None
 
     def proxy(self, asset, system_user):
-        try:
         self.server = self.get_server_conn(asset, system_user)
-        except PermissionFailed:
-            self.client.send(b"No permission")
+        if self.server is None:
             return
-
         session = Session(self.client, self.server)
         self.app.sessions.append(session)
         self.watch_win_size_change_async()
@@ -51,7 +47,8 @@ class ProxyServer:
 
     def get_server_conn(self, asset, system_user):
         if not self.validate_permission(asset, system_user):
-            raise PermissionFailed
+            self.client.send(b'No permission')
+            return None
 
         self.get_system_user_auth(system_user)
         ssh = paramiko.SSHClient()
@@ -62,12 +59,12 @@ class ProxyServer:
                         password=system_user.password,
                         pkey=system_user.private_key)
         except paramiko.AuthenticationException as e:
-            self.client.send(b"Authentication failed: %s" % e)
-            return
+            self.client.send("Authentication failed: {}".format(e).encode("utf-8"))
+            return None
 
         except socket.error as e:
-            self.client.send(b"Connection server error: %s" % e)
-            return
+            self.client.send("Connection server error: {}".format(e).encode("utf-8"))
+            return None
 
         term = self.request.meta.get('term', 'xterm')
         width = self.request.meta.get('width', 80)
@@ -88,4 +85,3 @@ class ProxyServer:
         thread.daemon = True
         thread.start()
 
-

coco/interactive.py

@@ -7,6 +7,7 @@ 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
+from .session import Session
 
 
 class InteractiveServer:
@@ -32,7 +33,7 @@ class InteractiveServer:
         0) 输入 \033[32mQ/q\033[0m 退出.\r\n""" % self.request.user
         self.client.send(banner.encode('utf-8'))
 
-    def get_choice(self, prompt='Opt> '):
+    def get_choice(self, prompt=b'Opt> '):
         """实现了一个ssh input, 提示用户输入, 获取并返回
 
         :return user input string
@@ -41,7 +42,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).encode('utf-8'))
+        self.client.send(wr(prompt, before=1, after=0))
         while True:
             data = self.client.recv(10)
             if len(data) == 0:
@@ -73,7 +74,8 @@ 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(b'', after=2))
-                option = parser.parse_input(b''.join(input_data))
+                option = parser.parse_input(input_data)
+                del input_data[:]
                 return option.strip()
             else:
                 self.client.send(data)
@@ -124,14 +126,19 @@ class InteractiveServer:
         pass
 
     def search_and_proxy(self, opt, from_result=False):
-        asset = Asset(id=1, hostname="testserver", ip="123.57.183.135", port=8022)
-        system_user = SystemUser(id=2, username="web", password="redhat123", name="web")
+        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)
 
     def connect(self, asset, system_user):
         forwarder = ProxyServer(self.app, self.client, self.request)
         forwarder.proxy(asset, system_user)
 
+    def replay_session(self, session_id):
+        session = Session(self.client, None)
+        session.id = "5a5dbfbe-093f-4bc1-810f-e8401b9e6045"
+        session.replay()
+
     def activate(self):
         self.display_banner()
         while True:

coco/interface.py

@@ -84,7 +84,7 @@ class SSHInterface(paramiko.ServerInterface):
     def check_channel_pty_request(
             self, channel, term, width, height,
             pixelwidth, pixelheight, modes):
-        logger.debug("Check channel pty request: %s %s %s %s %s" %
+        logger.info("Check channel pty request: %s %s %s %s %s" %
                      (term, width, height, pixelwidth, pixelheight))
         self.request.type = 'pty'
         self.request.meta = {
@@ -101,8 +101,6 @@ class SSHInterface(paramiko.ServerInterface):
 
     def check_channel_shell_request(self, channel):
         logger.info("Check channel shell request: %s" % channel)
-        self.request.type = 'shell'
-        self.request.meta = {'channel': channel}
         self.event.set()
         return True
 

coco/models.py

@@ -2,6 +2,8 @@ import json
 import threading
 import datetime
 
+from . import char
+from . import utils
 
 BUF_SIZE = 4096
 
@@ -108,20 +110,74 @@ class Server:
         self.chan = chan
         self.asset = asset
         self.system_user = system_user
+        self.send_bytes = 0
+        self.recv_bytes = 0
+
+        self.input_data = []
+        self.output_data = []
+        self.input = ''
+        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):
         return self.chan.fileno()
 
     def send(self, b):
+        if not self._input_initial:
+            self._input_initial = True
+
+        if self._have_enter_char(b):
+            self._in_input_state = False
+        else:
+            if not self._in_input_state:
+                print("#" * 30 + " 新周期 " + "#" * 30)
+                self._parse_input()
+                self._parse_output()
+                del self.input_data[:]
+                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)
 
     def recv(self, size):
-        return self.chan.recv(size)
+        data = self.chan.recv(size)
+        print("Recv: %s" % data)
+        if self._input_initial:
+            if self._in_input_state:
+                self.input_data.append(data)
+            else:
+                self.output_data.append(data)
+
+        return data
 
     def close(self):
         self.chan.close()
         return self.chan.transport.close()
 
+    @staticmethod
+    def _have_enter_char(s):
+        for c in char.ENTER_CHAR:
+            if c in s:
+                return True
+        return False
+
+    def _parse_output(self):
+        parser = utils.TtyIOParser()
+        print("\tOutput: %s" % parser.parse_output(self.output_data))
+
+    def _parse_input(self):
+        parser = utils.TtyIOParser()
+        print("\tInput: %s" % parser.parse_input(self.input_data))
+
     def __getattr__(self, item):
         return getattr(self.chan, item)
 

coco/session.py

@@ -24,6 +24,7 @@ class Session:
         self.watchers = []  # Only watch session
         self.sharers = []  # Join to the session, read and write
         self.running = True
+        self.replaying = True
         self.date_created = datetime.datetime.now()
         self.date_finished = None
         self.sel = selectors.DefaultSelector()
@@ -111,11 +112,12 @@ class Session:
         Record the session to a file. Using it replay in the future
         :return:
         """
+        logger.info("Start record session %s" % self.id)
         parent, child = socket.socketpair()
         self.add_watcher(parent)
-        with open(os.path.join(self.record_dir, self.id + ".rec"), 'wb') as screenf, \
+        with open(os.path.join(self.record_dir, self.id + ".rec"), 'wb') as dataf, \
             open(os.path.join(self.record_dir, self.id + ".time"), "w") as timef:
-            screenf.write("Script started on {}\n".format(time.asctime()).encode("utf-8"))
+            dataf.write("Script started on {}\n".format(time.asctime()).encode("utf-8"))
 
             while self.running:
                 start_t = time.time()
@@ -125,10 +127,9 @@ class Session:
                 if size == 0:
                     break
                 timef.write("%.4f %s\n" % (end_t-start_t, size))
-                screenf.write(data)
-                print("Pass %.4f, print %d" % (end_t-start_t, size))
-                print("Data: {}".format(data.decode('utf-8')))
-            screenf.write("Script done on {}\n".format(time.asctime()).encode("utf-8"))
+                dataf.write(data)
+            dataf.write("Script done on {}\n".format(time.asctime()).encode("utf-8"))
+        logger.info("End session record %s" % self.id)
 
     def record_async(self):
         thread = threading.Thread(target=self.record)
@@ -140,9 +141,26 @@ class Session:
         Replay the session
         :return:
         """
-        pass
+        with open(os.path.join(self.record_dir, self.id + ".rec"), 'rb') as dataf, \
+               open(os.path.join(self.record_dir, self.id + ".time"), "r") as timef:
 
-    def replay_down(self):
+            self.client.send(dataf.readline())
+            for l in timef:
+                if not self.replaying:
+                    break
+                t, size = float(l.split()[0]), int(l.split()[1])
+                data = dataf.read(size)
+                print("Sleep %s and send %s" % (t, data))
+                time.sleep(t)
+                self.client.send(data)
+        self.client.send("Replay session {} end".format(self.id).encode('utf-8'))
+        self.replaying = False
+
+    def replay_download(self):
+        """
+        Using termrecord generate html, then down user download it and share it
+        :return:
+        """
         pass
 
     def close(self):
@@ -152,7 +170,9 @@ class Session:
 
     def __str__(self):
         return self.id
-    __repr__ = __str__
+
+    def __repr__(self):
+        return self.id
 
 
 

coco/utils.py

@@ -138,11 +138,17 @@ class TtyIOParser(object):
         return self.ps1_pattern.sub('', command)
 
     def parse_output(self, data, sep='\n'):
+        """
+        Parse user command output
+
+        :param data: output data list like, [b'data', b'data']
+        :param sep:  line separator
+        :return: output unicode data
+        """
         output = []
-        if not isinstance(data, bytes):
-            data = data.encode('utf-8', 'ignore')
 
-        self.stream.feed(data)
+        for d in data:
+            self.stream.feed(d)
         for line in self.screen.display:
             if line.strip():
                 output.append(line)
@@ -150,11 +156,15 @@ class TtyIOParser(object):
         return sep.join(output[0:-1])
 
     def parse_input(self, data):
-        command = []
-        if not isinstance(data, bytes):
-            data = data.encode('utf-8', 'ignore')
+        """
+        Parse user input command
 
-        self.stream.feed(data)
+        :param data: input data list, like [b'data', b'data']
+        :return: command unicode
+        """
+        command = []
+        for d in data:
+            self.stream.feed(d)
         for line in self.screen.display:
             line = line.strip()
             if line:

conf_example.py

@@ -51,4 +51,5 @@ APP_NAME = "coco"
 # 和Jumpserver 保持心跳时间间隔
 # HEARTBEAT_INTERVAL = 5
 
-
+# 异步上报统计命令
+# COMMAND_PUSH_ASYNC = True