Merged in dev (pull request #27)

Dev

coco/app.py

@@ -205,6 +205,7 @@ class Coco:
         for client in self.clients:
             self.remove_client(client)
         time.sleep(1)
+        self.heartbeat()
         self.stop_evt.set()
         self.sshd.shutdown()
         self.httpd.shutdown()
@@ -234,4 +235,4 @@ class Coco:
         with self.lock:
             logger.info("Remove session: {}".format(session))
             self.sessions.remove(session)
-            self.service.finish_session(session.id)
+            self.service.finish_session(session.to_json())

coco/char.py

@@ -4,7 +4,9 @@
 
 BACKSPACE_CHAR = {b'\x08': b'\x08\x1b[K', b'\x7f': b'\x08\x1b[K'}
 ENTER_CHAR = [b'\r', b'\n', b'\r\n']
+ENTER_CHAR_ORDER = [ord(b'\r'), ord(b'\n')]
 UNSUPPORTED_CHAR = {b'\x15': 'Ctrl-U', b'\x0c': 'Ctrl-L', b'\x05': 'Ctrl-E'}
 CLEAR_CHAR = b'\x1b[H\x1b[2J'
 BELL_CHAR = b'\x07'
 NEW_LINE = b'\r\n'
+RZ_PROTOCOL_CHAR = b'**\x18B0900000000a87c\r\x8a\x11'

coco/httpd.py

@@ -11,7 +11,7 @@ from .models import Request, Client, WSProxy
 from .proxy import ProxyServer
 from .utils import get_logger
 
-__version__ = '0.4.0'
+__version__ = '0.5.0'
 BASE_DIR = os.path.dirname(os.path.dirname(__file__))
 
 logger = get_logger(__file__)

coco/interactive.py

@@ -80,6 +80,7 @@ class InteractiveServer:
         input_data = []
         parser = TtyIOParser()
         self.client.send(wr(prompt, before=1, after=0))
+
         while True:
             data = self.client.recv(10)
             logger.debug(data)
@@ -113,7 +114,7 @@ class InteractiveServer:
 
             # handle shell expect
             multi_char_with_enter = False
-            if len(data) > 1 and data[-1] in char.ENTER_CHAR:
+            if len(data) > 1 and data[-1] in char.ENTER_CHAR_ORDER:
                 self.client.send(data)
                 input_data.append(data[:-1])
                 multi_char_with_enter = True

coco/models.py

@@ -22,6 +22,18 @@ class Request:
         self.date_start = datetime.datetime.now()
 
 
+class SizedList(list):
+    def __init__(self, maxsize=0):
+        self.maxsize = maxsize
+        self.size = 0
+        super().__init__()
+
+    def append(self, b):
+        if self.maxsize == 0 or self.size < self.maxsize:
+            super().append(b)
+            self.size += len(b)
+
+
 class Client:
     """
     Client is the request client. Nothing more to say
@@ -78,8 +90,8 @@ class Server:
         self.recv_bytes = 0
         self.stop_evt = threading.Event()
 
-        self.input_data = []
-        self.output_data = []
+        self.input_data = SizedList(maxsize=1024)
+        self.output_data = SizedList(maxsize=1024)
         self._in_input_state = True
         self._input_initial = False
         self._in_vim_state = False
@@ -118,7 +130,8 @@ class Server:
                     self._input, self._output,
                     "#" * 30 + " End " + "#" * 30,
                 ))
-                self.session.put_command(self._input, self._output)
+                if self._input:
+                    self.session.put_command(self._input, self._output)
                 del self.input_data[:]
                 del self.output_data[:]
             self._in_input_state = True
@@ -152,10 +165,14 @@ class Server:
         return False
 
     def _parse_output(self):
+        if not self.output_data:
+            return ''
         parser = utils.TtyIOParser()
         return parser.parse_output(self.output_data)
 
     def _parse_input(self):
+        if not self.input_data or self.input_data[0] == char.RZ_PROTOCOL_CHAR:
+            return
         parser = utils.TtyIOParser()
         return parser.parse_input(self.input_data)
 
@@ -234,3 +251,7 @@ class WSProxy:
         self.child.close()
         self.ws.logout(self.connection)
         logger.debug("Proxy {} closed".format(self))
+
+
+
+

coco/recorder.py

@@ -129,11 +129,45 @@ class ServerReplayRecorder(ReplayRecorder):
             logger.info("Succeed to push {}'s {}".format(session_id, "record"))
         else:
             logger.error("Failed to push {}'s {}".format(session_id, "record"))
+        self.push_to_server(session_id)
 
     def push_to_server(self, session_id):
+        if self.upload_replay(3, session_id):
+            if self.finish_replay(3, session_id):
+                return True
+            else:
+                return False
+        else:
+            return False
+
+    def push_local(self, session_id):
         return self.app.service.push_session_replay(os.path.join(self.app.config['LOG_DIR'], session_id + '.replay.gz'),
                                                     session_id)
 
+    def upload_replay(self, times, session_id):
+        if times > 0:
+            if self.push_local(session_id):
+                logger.info("success push session: {}'s replay log ".format(session_id))
+                return True
+            else:
+                logger.error("failed report session {}'s replay log, try  {} times".format(session_id, times))
+                return self.upload_replay(times - 1, session_id)
+        else:
+            logger.error("failed report session {}'s replay log".format(session_id))
+            return False
+
+    def finish_replay(self, times, session_id):
+        if times > 0:
+            if self.app.service.finish_replay(session_id):
+                logger.info("success report session {}'s replay log ".format(session_id))
+                return True
+            else:
+                logger.error("failed report session {}'s replay log, try {} times".format(session_id, times))
+                return self.finish_replay(times - 1, session_id)
+        else:
+            logger.error("failed report session {}'s replay log".format(session_id))
+            return False
+
     def __del__(self):
         print("{} has been gc".format(self))
         del self.file
@@ -254,18 +288,41 @@ class S3ReplayRecorder(ServerReplayRecorder):
         else:
             self.s3 = boto3.client('s3')
 
-    def push_to_server(self, session_id):
+    def push_to_s3(self, session_id):
         logger.debug("push to server")
         try:
             self.s3.upload_file(
                 os.path.join(self.app.config['LOG_DIR'], session_id + '.replay.gz'),
                 self.bucket,
-                self.app.config.get("NAME", "coco") + time.strftime('%Y-%m-%d', time.localtime(
+                time.strftime('%Y-%m-%d', time.localtime(
                     self.starttime)) + '/' + session_id + '.replay.gz')
+            return True
         except:
-            return self.app.service.push_session_replay(
-                os.path.join(self.app.config['LOG_DIR'], session_id + '.replay.gz'),
-                session_id)
+            return False
+
+    def upload_replay(self, times, session_id):
+        if times > 0:
+            if self.push_to_s3(session_id):
+                logger.info("success push session: {}'s replay log to S3 ".format(session_id))
+                return True
+            else:
+                logger.error("failed report session {}'s replay log to S3, try  {} times".format(session_id, times))
+                return self.upload_replay(times - 1, session_id)
+        else:
+            logger.error("failed report session {}'s replay log S3, try to push to local".format(session_id))
+            return self.upload_replay_to_local(3, session_id)
+
+    def upload_replay_to_local(self, times, session_id):
+        if times > 0:
+            if self.push_local(session_id):
+                logger.info("success push session: {}'s replay log ".format(session_id))
+                return True
+            else:
+                logger.error("failed report session {}'s replay log, try  {} times".format(session_id, times))
+                return self.upload_replay_to_local(times - 1, session_id)
+        else:
+            logger.error("failed report session {}'s replay log".format(session_id))
+            return False
 
 
 def get_command_recorder_class(config):
@@ -280,6 +337,7 @@ def get_command_recorder_class(config):
 
 def get_replay_recorder_class(config):
     replay_storage = config["REPLAY_STORAGE"]
+    logger.debug(replay_storage)
     storage_type = replay_storage.get('TYPE')
     if storage_type == "s3":
         return S3ReplayRecorder