[Feture] 拆分replay和command recorder

coco/app.py

@@ -25,7 +25,7 @@ class Coco:
         'DEBUG': True,
         'BIND_HOST': '0.0.0.0',
         'SSHD_PORT': 2222,
-        'WS_PORT': 5000,
+        'HTTPD_PORT': 5000,
         'ACCESS_KEY': '',
         'ACCESS_KEY_ENV': 'COCO_ACCESS_KEY',
         'ACCESS_KEY_FILE': os.path.join(BASE_DIR, 'keys', '.access_key'),
@@ -33,12 +33,14 @@ class Coco:
         'LOG_LEVEL': 'INFO',
         'LOG_DIR': os.path.join(BASE_DIR, 'logs'),
         'SESSION_DIR': os.path.join(BASE_DIR, 'sessions'),
-        'SESSION_COMMAND_STORE': "server",  # elasticsearch
+        'REPLAY_STORE_ENGINE': 'server',   # local, server
+        'COMMAND_STORE_ENGINE': 'server',  # local, server, elasticsearch(not yet)
         'ASSET_LIST_SORT_BY': 'hostname',  # hostname, ip
         'SSH_PASSWORD_AUTH': True,
         'SSH_PUBLIC_KEY_AUTH': True,
         'HEARTBEAT_INTERVAL': 5,
         'MAX_CONNECTIONS': 500,
+        # 'MAX_RECORD_OUTPUT_LENGTH': 4096,
     }
 
     def __init__(self, name=None, root_path=None):
@@ -130,7 +132,7 @@ class Coco:
             if self.config["SSHD_PORT"] != 0:
                 self.run_sshd()
 
-            if self.config['WS_PORT'] != 0:
+            if self.config['HTTPD_PORT'] != 0:
                 self.run_httpd()
 
             self.stop_evt.wait()

coco/forward.py

@@ -8,7 +8,7 @@ import paramiko
 
 from .session import Session
 from .models import Server
-from .record import FileRecorder
+from .record import LocalFileReplayRecorder, LocalFileCommandRecorder
 from .utils import wrap_with_line_feed as wr
 
 
@@ -33,10 +33,11 @@ class ProxyServer:
         session = Session(self.client, self.server)
         self.app.add_session(session)
         self.watch_win_size_change_async()
-        recorder = FileRecorder(self.app, session)
-        session.add_recorder(recorder)
+        replay_recorder = LocalFileReplayRecorder(self.app, session)
+        session.add_recorder(replay_recorder)
         session.record_replay_async()
-        self.server.add_recorder(recorder)
+        cmd_recorder = LocalFileCommandRecorder(self.app, session)
+        self.server.add_recorder(cmd_recorder)
         self.server.record_command_async()
         session.bridge()
         session.stop_evt.set()

coco/httpd.py

@@ -98,7 +98,7 @@ class HttpServer:
 
     def run(self):
         host = self.app.config["BIND_HOST"]
-        port = self.app.config["WS_PORT"]
+        port = self.app.config["HTTPD_PORT"]
         print('Starting websocket server at %(host)s:%(port)s' %
               {"host": host, "port": port})
         ws = tornado.web.Application(self.routers, **self.settings)

coco/record.py

@@ -2,6 +2,7 @@
 #
 
 import abc
+import threading
 import time
 import os
 import logging
@@ -11,7 +12,7 @@ logger = logging.getLogger(__file__)
 BUF_SIZE = 1024
 
 
-class Recorder(metaclass=abc.ABCMeta):
+class ReplayRecorder(metaclass=abc.ABCMeta):
 
     def __init__(self, app, session):
         self.app = app
@@ -21,6 +22,20 @@ class Recorder(metaclass=abc.ABCMeta):
     def record_replay(self, now, timedelta, size, data):
         pass
 
+    @abc.abstractmethod
+    def start(self):
+        pass
+
+    @abc.abstractmethod
+    def done(self):
+        pass
+
+
+class CommandRecorder(metaclass=abc.ABCMeta):
+    def __init__(self, app, session):
+        self.app = app
+        self.session = session
+
     @abc.abstractmethod
     def record_command(self, now, _input, _output):
         pass
@@ -34,13 +49,12 @@ class Recorder(metaclass=abc.ABCMeta):
         pass
 
 
-class FileRecorder(Recorder):
+class LocalFileReplayRecorder(ReplayRecorder):
 
     def __init__(self, app, session):
         super().__init__(app, session)
         self.data_f = None
         self.time_f = None
-        self.cmd_f = None
         self.prepare_file()
 
     def prepare_file(self):
@@ -54,12 +68,10 @@ class FileRecorder(Recorder):
         filename = os.path.join(session_dir, str(self.session.id))
         data_filename = filename + ".rec"
         time_filename = filename + ".time"
-        cmd_filename = filename + ".cmd"
 
         try:
             self.data_f = open(data_filename, "wb")
             self.time_f = open(time_filename, "w")
-            self.cmd_f = open(cmd_filename, "w")
         except IOError as e:
             logger.debug(e)
             self.done()
@@ -69,13 +81,6 @@ class FileRecorder(Recorder):
         self.time_f.write("{} {}\n".format(timedelta, size))
         self.data_f.write(data)
 
-    def record_command(self, now, _input, _output):
-        logger.debug("File recorder command: ({},{})".format(_input, _output))
-        self.cmd_f.write("{}\n".format(now.strftime("%Y-%m-%d %H:%M:%S")))
-        self.cmd_f.write("$ {}\n".format(_input))
-        self.cmd_f.write("{}\n\n".format(_output))
-        self.cmd_f.flush()
-
     def start(self):
         logger.debug("Session {} start".format(self.session.id))
         self.data_f.write("Session {} started on {}\n".format(self.session.id, time.asctime()).encode("utf-8"))
@@ -83,12 +88,68 @@ class FileRecorder(Recorder):
     def done(self):
         logger.debug("Session {} record done".format(self.session.id))
         self.data_f.write("Session {} done on {}\n".format(self.session.id, time.asctime()).encode("utf-8"))
-        for f in [self.data_f, self.time_f, self.cmd_f]:
+        for f in (self.data_f, self.time_f):
             try:
                 f.close()
             except IOError:
                 pass
 
 
+class LocalFileCommandRecorder(CommandRecorder):
+    def __init__(self, app, session):
+        super().__init__(app, session)
+        self.cmd_f = None
+        self.prepare_file()
+
+    def prepare_file(self):
+        session_dir = os.path.join(
+            self.app.config["SESSION_DIR"],
+            self.session.date_created.strftime("%Y-%m-%d")
+        )
+        if not os.path.isdir(session_dir):
+            os.mkdir(session_dir)
+
+        filename = os.path.join(session_dir, str(self.session.id))
+        cmd_filename = filename + ".cmd"
+
+        try:
+            self.cmd_f = open(cmd_filename, "w")
+        except IOError as e:
+            logger.debug(e)
+            self.done()
+
+    def record_command(self, now, _input, _output):
+        logger.debug("File recorder command: ({},{})".format(_input, _output))
+        self.cmd_f.write("{}\n".format(now.strftime("%Y-%m-%d %H:%M:%S")))
+        self.cmd_f.write("$ {}\n".format(_input))
+        self.cmd_f.write("{}\n\n".format(_output))
+        self.cmd_f.flush()
+
+    def start(self):
+        pass
+
+    def done(self):
+        pass
+
+
+class ServerReplayRecorder(LocalFileReplayRecorder):
+
+    def done(self):
+        super().done()
+        self.push_records()
+
+    def push_records(self):
+        def func():
+            self.push_replay_record()
+
+        thread = threading.Thread(target=func)
+        thread.start()
+
+    def push_replay_record(self):
+        pass
+
+
+
+
 
 

conf_example.py

@@ -14,8 +14,8 @@ APP_NAME = "coco"
 # 监听的SSH端口号, 默认2222
 # SSHD_PORT = 2222
 
-# 监听的WS端口号，默认5000
-# WS_PORT = 5000
+# 监听的HTTP/WS端口号，默认5000
+# HTTPD_PORT = 5000
 
 # 是否开启DEBUG
 # DEBUG = True