Merged in test (pull request #11)

Test

coco/app.py

@@ -6,8 +6,8 @@ import datetime
 import os
 import time
 import threading
-import logging
 import socket
+import json
 
 from jms.service import AppService
 
@@ -17,12 +17,13 @@ from .httpd import HttpServer
 from .logger import create_logger
 from .tasks import TaskHandler
 from .recorder import get_command_recorder_class, get_replay_recorder_class
+from .utils import get_logger
 
 
-__version__ = '0.4.0'
+__version__ = '0.5.0'
 
 BASE_DIR = os.path.dirname(os.path.dirname(__file__))
-logger = logging.getLogger(__file__)
+logger = get_logger(__file__)
 
 
 class Coco:
@@ -42,19 +43,18 @@ class Coco:
         'LOG_DIR': os.path.join(BASE_DIR, 'logs'),
         'SESSION_DIR': os.path.join(BASE_DIR, 'sessions'),
         'ASSET_LIST_SORT_BY': 'hostname',  # hostname, ip
-        'SSH_PASSWORD_AUTH': True,
-        'SSH_PUBLIC_KEY_AUTH': True,
+        'PASSWORD_AUTH': True,
+        'PUBLIC_KEY_AUTH': True,
         'HEARTBEAT_INTERVAL': 5,
         'MAX_CONNECTIONS': 500,
         'ADMINS': '',
-        'REPLAY_RECORD_ENGINE': 'server',   # local, server
-        'COMMAND_RECORD_ENGINE': 'server',  # local, server, elasticsearch(not yet)
+        'COMMAND_STORAGE': {'TYPE': 'server'},   # server
+        'REPLAY_RECORD_ENGINE': 'server',
     }
 
-    def __init__(self, name=None, root_path=None):
+    def __init__(self, root_path=None):
         self.root_path = root_path if root_path else BASE_DIR
         self.config = self.config_class(self.root_path, defaults=self.default_config)
-        self.name = name if name else self.config["NAME"]
         self.sessions = []
         self.clients = []
         self.lock = threading.Lock()
@@ -66,6 +66,10 @@ class Coco:
         self.command_recorder_class = None
         self._task_handler = None
 
+    @property
+    def name(self):
+        return self.config["NAME"]
+
     @property
     def service(self):
         if self._service is None:
@@ -93,16 +97,20 @@ class Coco:
     def make_logger(self):
         create_logger(self)
 
-    # Todo: load some config from server like replay and common upload
     def load_extra_conf_from_server(self):
-        pass
+        configs = self.service.load_config_from_server()
+        logger.debug("Loading config from server: {}".format(
+            json.dumps(configs)
+        ))
+        self.config.update(configs)
 
-    def initial_recorder(self):
-        self.replay_recorder_class = get_replay_recorder_class(self)
-        self.command_recorder_class = get_command_recorder_class(self)
+    def get_recorder_class(self):
+        self.replay_recorder_class = get_replay_recorder_class(self.config)
+        self.command_recorder_class = get_command_recorder_class(self.config)
 
     def new_command_recorder(self):
-        return self.command_recorder_class(self)
+        recorder = self.command_recorder_class(self)
+        return recorder
 
     def new_replay_recorder(self):
         return self.replay_recorder_class(self)
@@ -111,7 +119,7 @@ class Coco:
         self.make_logger()
         self.service.initial()
         self.load_extra_conf_from_server()
-        self.initial_recorder()
+        self.get_recorder_class()
         self.keep_heartbeat()
         self.monitor_sessions()
 

coco/httpd.py

@@ -2,7 +2,6 @@
 # -*- coding: utf-8 -*-
 #
 import os
-import logging
 import socket
 from flask_socketio import SocketIO, Namespace, emit, join_room, leave_room
 from flask import Flask, send_from_directory, render_template, request, jsonify
@@ -12,11 +11,12 @@ import uuid
 from jms.models import User
 from .models import Request, Client, WSProxy
 from .proxy import ProxyServer
+from .utils import get_logger
 
 __version__ = '0.4.0'
 BASE_DIR = os.path.dirname(os.path.dirname(__file__))
 
-logger = logging.getLogger(__file__)
+logger = get_logger(__file__)
 
 
 class BaseWebSocketHandler:

coco/interactive.py

@@ -2,22 +2,21 @@
 # -*- coding: utf-8 -*-
 #
 
-import logging
 import socket
 import threading
 import weakref
-
 import os
+
 from jms.models import Asset, AssetGroup
 
 from . import char
 from .utils import wrap_with_line_feed as wr, wrap_with_title as title, \
     wrap_with_primary as primary, wrap_with_warning as warning, \
     is_obj_attr_has, is_obj_attr_eq, sort_assets, TtyIOParser, \
-    ugettext as _
+    ugettext as _, get_logger
 from .proxy import ProxyServer
 
-logger = logging.getLogger(__file__)
+logger = get_logger(__file__)
 
 
 class InteractiveServer:
@@ -42,7 +41,7 @@ class InteractiveServer:
         if self._search_result:
             return self._search_result
         else:
-            return None
+            return []
 
     @search_result.setter
     def search_result(self, value):
@@ -272,7 +271,7 @@ class InteractiveServer:
 
     def search_and_proxy(self, opt):
         self.search_assets(opt)
-        if len(self.search_result) == 1:
+        if self.search_result and len(self.search_result) == 1:
             self.proxy(self.search_result[0])
         else:
             self.display_search_result()

coco/interface.py

@@ -2,13 +2,13 @@
 # -*- coding: utf-8 -*-
 #
 
-import logging
 import paramiko
 import threading
 import weakref
 
+from .utils import get_logger
 
-logger = logging.getLogger(__file__)
+logger = get_logger(__file__)
 
 
 class SSHInterface(paramiko.ServerInterface):
@@ -43,9 +43,9 @@ class SSHInterface(paramiko.ServerInterface):
 
     def get_allowed_auths(self, username):
         supported = []
-        if self.app.config["SSH_PASSWORD_AUTH"]:
+        if self.app.config["PASSWORD_AUTH"]:
             supported.append("password")
-        if self.app.config["SSH_PUBLIC_KEY_AUTH"]:
+        if self.app.config["PUBLIC_KEY_AUTH"]:
             supported.append("publickey")
 
         return ",".join(supported)

coco/logger.py

@@ -4,43 +4,56 @@
 
 import os
 import logging
-from logging import StreamHandler
-from logging.handlers import TimedRotatingFileHandler
-
-
-LOG_LEVELS = {
-    'DEBUG': logging.DEBUG,
-    'INFO': logging.INFO,
-    'WARN': logging.WARNING,
-    'WARNING': logging.WARNING,
-    'ERROR': logging.ERROR,
-    'FATAL': logging.FATAL,
-    'CRITICAL': logging.CRITICAL,
-}
+from logging.config import dictConfig
 
 
 def create_logger(app):
     level = app.config['LOG_LEVEL']
-    level = LOG_LEVELS.get(level, logging.INFO)
     log_dir = app.config.get('LOG_DIR')
     log_path = os.path.join(log_dir, 'coco.log')
+    main_setting = {
+        'handlers': ['console', 'file'],
+        'level': level,
+        'propagate': False,
+    }
+    config = dict(
+        version=1,
+        formatters={
+            "main": {
+                'format': '%(asctime)s [%(module)s %(levelname)s] %(message)s',
+                'datefmt': '%Y-%m-%d %H:%M:%S',
+            },
+            'simple': {
+                'format': '%(asctime)s [%(levelname)-8s] %(message)s',
+                'datefmt': '%Y-%m-%d %H:%M:%S',
+            }
+        },
+        handlers={
+            'null': {
+                'level': 'DEBUG',
+                'class': 'logging.NullHandler',
+            },
+            'console': {
+                'level': 'DEBUG',
+                'class': 'logging.StreamHandler',
+                'formatter': 'main'
+            },
+            'file': {
+                'level': 'DEBUG',
+                'class': 'logging.FileHandler',
+                'formatter': 'main',
+                'filename': log_path,
+            },
+        },
+        loggers={
+            'coco': main_setting,
+            'paramiko': main_setting,
+            'jms': main_setting,
+        }
+    )
+
+    dictConfig(config)
     logger = logging.getLogger()
+    return logger
 
-    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 [%(levelname)s] %(message)s',
-    #     datefmt='%Y-%m-%d %H:%M:%S'
-    # )
-    console_handler = StreamHandler()
-    file_handler = TimedRotatingFileHandler(
-        filename=log_path, when='D', backupCount=10
-    )
 
-    for handler in [console_handler, file_handler]:
-        handler.setFormatter(main_formatter)
-        logger.addHandler(handler)
-    logger.setLevel(level)
-    logging.getLogger("requests").setLevel(logging.WARNING)

coco/models.py

@@ -2,14 +2,13 @@
 # -*- coding: utf-8 -*-
 import threading
 import datetime
-import logging
 import weakref
 
 from . import char
 from . import utils
 
 BUF_SIZE = 4096
-logger = logging.getLogger(__file__)
+logger = utils.get_logger(__file__)
 
 
 class Request:

coco/proxy.py

@@ -4,7 +4,6 @@
 
 import socket
 import threading
-import logging
 import time
 import weakref
 import paramiko
@@ -13,10 +12,10 @@ from paramiko.ssh_exception import SSHException
 from .session import Session
 from .models import Server
 from .utils import wrap_with_line_feed as wr, wrap_with_warning as warning, \
-    get_private_key_fingerprint
+    get_private_key_fingerprint, get_logger
 
 
-logger = logging.getLogger(__file__)
+logger = get_logger(__file__)
 TIMEOUT = 8
 BUF_SIZE = 4096
 

coco/recorder.py

@@ -3,7 +3,6 @@
 #
 
 import abc
-import logging
 import threading
 import time
 import os
@@ -11,9 +10,12 @@ import gzip
 import json
 import shutil
 
+from jms_es_sdk import ESStore
+
+from .utils import get_logger
 from .alignment import MemoryQueue
 
-logger = logging.getLogger(__file__)
+logger = get_logger(__file__)
 BUF_SIZE = 1024
 
 
@@ -183,17 +185,70 @@ class ServerCommandRecorder(CommandRecorder, metaclass=Singleton):
         print("{} has been gc".format(self))
 
 
-def get_command_recorder_class(app):
-    command_engine = app.config["COMMAND_RECORD_ENGINE"]
+class ESCommandRecorder(CommandRecorder, metaclass=Singleton):
+    batch_size = 10
+    timeout = 5
+    no = 0
+    default_hosts = ["http://localhost"]
 
-    if command_engine == "server":
-        return ServerCommandRecorder
+    def __init__(self, app):
+        super().__init__(app)
+        self.queue = MemoryQueue()
+        self.stop_evt = threading.Event()
+        self.push_to_es_async()
+        self.__class__.no += 1
+        self.store = ESStore(app.config["COMMAND_STORAGE"].get("HOSTS", self.default_hosts))
+        if not self.store.ping():
+            raise AssertionError("ESCommand storage init error")
+
+    def record(self, data):
+        if data and data['input']:
+            data['input'] = data['input'][:128]
+            data['output'] = data['output'][:1024]
+            data['timestamp'] = int(data['timestamp'])
+            self.queue.put(data)
+
+    def push_to_es_async(self):
+        def func():
+            while not self.stop_evt.is_set():
+                data_set = self.queue.mget(self.batch_size,
+                                           timeout=self.timeout)
+                logger.debug(
+                    "<Session command recorder {}> queue size: {}".format(
+                        self.no, self.queue.qsize())
+                )
+                if not data_set:
+                    continue
+                logger.debug("Send {} commands to server".format(len(data_set)))
+                ok = self.store.bulk_save(data_set)
+                if not ok:
+                    self.queue.mput(data_set)
+
+        thread = threading.Thread(target=func)
+        thread.daemon = True
+        thread.start()
+
+    def session_start(self, session_id):
+        pass
+
+    def session_end(self, session_id):
+        pass
+
+    def __del__(self):
+        print("{} has been gc".format(self))
+
+
+def get_command_recorder_class(config):
+    command_storage = config["COMMAND_STORAGE"]
+
+    if command_storage['TYPE'] == "elasticsearch":
+        return ESCommandRecorder
     else:
         return ServerCommandRecorder
 
 
-def get_replay_recorder_class(app):
-    replay_engine = app.config["REPLAY_RECORD_ENGINE"]
+def get_replay_recorder_class(config):
+    replay_engine = config["REPLAY_RECORD_ENGINE"]
     if replay_engine == "server":
         return ServerReplayRecorder
     else:

coco/session.py

@@ -3,13 +3,14 @@
 #
 import threading
 import uuid
-import logging
 import datetime
 import selectors
 import time
 
+from .utils import get_logger
+
 BUF_SIZE = 1024
-logger = logging.getLogger(__file__)
+logger = get_logger(__file__)
 
 
 class Session:

coco/sshd.py

@@ -3,20 +3,16 @@
 #
 
 import os
-import logging
 import socket
 import threading
 import paramiko
-import sys
 
-import time
-
-from .utils import ssh_key_gen
+from .utils import ssh_key_gen, get_logger
 from .interface import SSHInterface
 from .interactive import InteractiveServer
 from .models import Client, Request
 
-logger = logging.getLogger(__file__)
+logger = get_logger(__file__)
 BACKLOG = 5
 
 
@@ -90,14 +86,12 @@ class SSHServer:
 
     def handle_chan(self, chan, request):
         client = Client(chan, request)
-        print(chan)
-        print(request)
         self.app.add_client(client)
         self.dispatch(client)
 
     def dispatch(self, client):
         request_type = client.request.type
-        if request_type == 'pty':
+        if request_type == 'pty' or request_type == 'x11':
             logger.info("Request type `pty`, dispatch to interactive mode")
             InteractiveServer(self.app, client).interact()
         elif request_type == 'exec':

coco/tasks.py

@@ -2,9 +2,10 @@
 # -*- coding: utf-8 -*-
 #
 import weakref
-import logging
 
-logger = logging.getLogger(__file__)
+from .utils import get_logger
+
+logger = get_logger(__file__)
 
 
 class TaskHandler:

coco/utils.py

@@ -5,6 +5,7 @@
 from __future__ import unicode_literals
 
 import hashlib
+import logging
 import re
 import os
 import threading
@@ -371,4 +372,8 @@ def compile_message():
     pass
 
 
+def get_logger(file_name):
+    return logging.getLogger('coco.'+file_name)
+
+
 ugettext = _gettext()

conf_docker.py

@@ -60,6 +60,10 @@ class Config:
     # Admin的名字，出问题会提示给用户
     ADMINS = os.environ.get("ADMINS") or ''
 
+    COMMAND_STORAGE = {
+        "TYPE": "server"
+    }
+
 
 class ConfigDocker(Config):
     pass

conf_example.py

@@ -9,10 +9,10 @@ BASE_DIR = os.path.dirname(__file__)
 
 class Config:
     """
-    Coco config file
+    Coco config file, coco also load config from server update setting below
     """
     # 项目名称, 会用来向Jumpserver注册, 识别而已, 不能重复
-    # APP_NAME = "localhost"
+    # NAME = "localhost"
 
     # Jumpserver项目的url, api请求注册会使用
     # CORE_HOST = os.environ.get("CORE_HOST") or 'http://127.0.0.1:8080'
@@ -49,16 +49,19 @@ class Config:
     # ASSET_LIST_SORT_BY = 'ip'
 
     # 登录是否支持密码认证
-    # SSH_PASSWORD_AUTH = True
+    # PASSWORD_AUTH = True
 
     # 登录是否支持秘钥认证
-    # SSH_PUBLIC_KEY_AUTH = True
+    # PUBLIC_KEY_AUTH = True
 
     # 和Jumpserver 保持心跳时间间隔
     # HEARTBEAT_INTERVAL = 5
 
     # Admin的名字，出问题会提示给用户
     # ADMINS = ''
+    COMMAND_STORAGE = {
+        "TYPE": "server"
+    }
 
 
 config = Config()

requirements/requirements.txt

@@ -28,4 +28,5 @@ tornado==4.5.2
 urllib3==1.22
 wcwidth==0.1.7
 werkzeug==0.12.2
-jumpserver-python-sdk==0.0.23
+jumpserver-python-sdk==0.0.26
+jms-es-sdk