Skip to content

C
coco
Commit 5ec6eae5 authored by 广宏伟's avatar 广宏伟
Browse files
Options

Merged in dev (pull request #58) 

Dev
parents 2467b4e2 a2d2e325
Show whitespace changes
Inline Side-by-side
Showing with 94 additions and 184 deletions
coco/app.py
View file @ 5ec6eae5
......@@ -17,7 +17,7 @@ from .sshd import SSHServer
from .httpd import HttpServer
from .logger import create_logger
from .tasks import TaskHandler
from .recorder import get_command_recorder_class, ServerReplayRecorder
from .recorder import ReplayRecorder, CommandRecorder
from .utils import get_logger, register_app, register_service
......@@ -56,7 +56,6 @@ class Coco:
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.sessions = []
self.clients = []
self.lock = threading.Lock()
......@@ -67,8 +66,15 @@ class Coco:
self.replay_recorder_class = None
self.command_recorder_class = None
self._task_handler = None
self.config = None
self.init_config()
register_app(self)
def init_config(self):
self.config = self.config_class(
self.root_path, defaults=self.default_config
)
@property
def name(self):
if self.config['NAME']:
......@@ -111,23 +117,21 @@ class Coco:
))
self.config.update(configs)
def get_recorder_class(self):
self.replay_recorder_class = ServerReplayRecorder
self.command_recorder_class = get_command_recorder_class(self.config)
def new_command_recorder(self):
return self.command_recorder_class()
@staticmethod
def new_command_recorder():
return CommandRecorder()
def new_replay_recorder(self):
return self.replay_recorder_class()
@staticmethod
def new_replay_recorder():
return ReplayRecorder()
def bootstrap(self):
self.make_logger()
self.service.initial()
self.load_extra_conf_from_server()
self.get_recorder_class()
self.keep_heartbeat()
self.monitor_sessions()
self.monitor_sessions_replay()
def heartbeat(self):
_sessions = [s.to_json() for s in self.sessions]
......@@ -156,6 +160,26 @@ class Coco:
thread = threading.Thread(target=func)
thread.start()
def monitor_sessions_replay(self):
interval = 10
recorder = self.new_replay_recorder()
log_dir = os.path.join(self.config['LOG_DIR'])
def func():
while not self.stop_evt.is_set():
active_sessions = [str(session.id) for session in self.sessions]
for filename in os.listdir(log_dir):
session_id = filename.split('.')[0]
if len(session_id) != 36:
continue
if session_id not in active_sessions:
recorder.file_path = os.path.join(log_dir, filename)
recorder.upload_replay(session_id, 1)
time.sleep(interval)
thread = threading.Thread(target=func)
thread.start()
def monitor_sessions(self):
interval = self.config["HEARTBEAT_INTERVAL"]
......
coco/httpd.py
View file @ 5ec6eae5
......@@ -4,7 +4,7 @@
import os
import socket
import uuid
import traceback
from copy import deepcopy
from flask_socketio import SocketIO, Namespace, join_room
from flask import Flask, request, current_app, redirect
......
coco/interactive.py
View file @ 5ec6eae5
......@@ -38,6 +38,9 @@ class InteractiveServer:
@search_result.setter
def search_result(self, value):
if not value:
self._search_result = value
return
value = self.filter_system_users(value)
self._search_result = value
......@@ -88,7 +91,7 @@ class InteractiveServer:
result = []
# 所有的
if q == '':
if q in ('', None):
result = self.assets
# 用户输入的是数字，可能想使用id唯一键搜索
elif q.isdigit() and self.search_result and \
......@@ -242,6 +245,7 @@ class InteractiveServer:
self.search_assets(opt)
if self.search_result and len(self.search_result) == 1:
asset = self.search_result[0]
self.search_result = None
if asset.platform == "Windows":
self.client.send(warning(
_("终端不支持登录windows, 请使用web terminal访问"))
......
coco/models.py
View file @ 5ec6eae5
......@@ -249,10 +249,12 @@ class WSProxy:
while not self.stop_event.is_set():
try:
data = self.child.recv(BUF_SIZE)
except OSError:
continue
except (OSError, EOFError):
self.close()
break
if len(data) == 0:
self.close()
break
data = data.decode(errors="ignore")
self.ws.emit("data", {'data': data, 'room': self.room_id},
room=self.room_id)
......@@ -265,6 +267,7 @@ class WSProxy:
thread.start()
def close(self):
self.ws.emit("logout", {"room": self.room_id}, room=self.room_id)
self.stop_event.set()
try:
self.child.shutdown(1)
......
coco/recorder.py
View file @ 5ec6eae5
......@@ -8,6 +8,7 @@ import time
import os
import gzip
import json
from copy import deepcopy
import jms_storage
......@@ -20,60 +21,6 @@ BUF_SIZE = 1024
class ReplayRecorder(metaclass=abc.ABCMeta):
def __init__(self, session=None):
self.session = session
@abc.abstractmethod
def record(self, data):
"""
记录replay数据
:param data: 数据 {
"session": "",
"data": "",
"timestamp": ""
}
:return:
"""
@abc.abstractmethod
def session_start(self, session_id):
print("Session start: {}".format(session_id))
pass
@abc.abstractmethod
def session_end(self, session_id):
print("Session end: {}".format(session_id))
pass
class CommandRecorder:
def __init__(self, session=None):
self.session = session
def record(self, data):
"""
:param data: 数据 {
"session":
"input":
"output":
"user":
"asset":
"system_user":
"timestamp":
}
:return:
"""
def session_start(self, session_id):
print("Session start: {}".format(session_id))
pass
def session_end(self, session_id):
print("Session end: {}".format(session_id))
pass
class ServerReplayRecorder(ReplayRecorder):
time_start = None
storage = None
......@@ -81,6 +28,7 @@ class ServerReplayRecorder(ReplayRecorder):
super().__init__()
self.file = None
self.file_path = None
self.get_storage()
def record(self, data):
"""
......@@ -107,70 +55,68 @@ class ServerReplayRecorder(ReplayRecorder):
def session_end(self, session_id):
self.file.write('"0":""}')
self.file.close()
if self.upload_replay(session_id):
logger.info("Succeed to push {}'s {}".format(session_id, "record"))
self.upload_replay(session_id)
def get_storage(self):
config = deepcopy(current_app.config["REPLAY_STORAGE"])
config["SERVICE"] = app_service
self.storage = jms_storage.get_object_storage(config)
def upload_replay(self, session_id, times=3):
if times < 1:
if self.storage.type == 'jms':
return False
else:
logger.error("Failed to push {}'s {}".format(session_id, "record"))
self.storage = jms_storage.JMSReplayStorage(app_service)
self.upload_replay(session_id, times=3)
def upload_replay(self, session_id):
configs = app_service.load_config_from_server()
logger.debug("upload_replay print config: {}".format(configs))
self.storage = jms_storage.init(configs["REPLAY_STORAGE"])
if not self.storage:
self.storage = jms_storage.jms(app_service)
if self.push_file(3, session_id):
ok, msg = self.push_to_storage(session_id)
if not ok:
msg = 'Failed push replay file: {}, try again {}'.format(msg, times)
logger.warn(msg)
self.upload_replay(session_id, times-1)
else:
msg = 'Success push replay file: {}'.format(session_id)
logger.info(msg)
self.finish_replay(3, session_id)
os.unlink(self.file_path)
return True
else:
return False
def push_to_storage(self, session_id):
dt = time.strftime('%Y-%m-%d', time.localtime(self.time_start))
target = dt + '/' + session_id + '.replay.gz'
return self.storage.upload_file(self.file_path, target)
def push_file(self, times, session_id):
if times < 0:
if self.storage.type() == 'jms':
return False
else:
msg = "Failed push session {}'s replay log to storage".format(session_id)
logger.error(msg)
self.storage = jms_storage.jms(app_service)
return self.push_file(3, session_id)
if self.push_to_storage(session_id):
logger.info("Success push session: {}'s replay log to storage ".format(session_id))
return True
else:
msg = "Failed push session {}'s replay log to storage, try {} times".format(session_id, times)
logger.error(msg)
return self.push_file(times - 1, session_id)
return self.storage.upload(self.file_path, target)
def finish_replay(self, times, session_id):
if times < 0:
logger.error("Failed finished session {}'s replay".format(session_id))
if times < 1:
logger.error(
"Failed finished session {}'s replay".format(session_id)
)
return False
if app_service.finish_replay(session_id):
logger.info("Success finish session {}'s replay ".format(session_id))
logger.info(
"Success finish session {}'s replay ".format(session_id)
)
return True
else:
logger.error("Failed finish session {}'s replay, try {} times".format(session_id, times))
msg = "Failed finish session {}'s replay, try {} times"
logger.error(msg.format(session_id, times))
return self.finish_replay(times - 1, session_id)
class ServerCommandRecorder(CommandRecorder, metaclass=Singleton):
class CommandRecorder(metaclass=Singleton):
batch_size = 10
timeout = 5
no = 0
storage = None
def __init__(self):
super().__init__()
self.queue = MemoryQueue()
self.stop_evt = threading.Event()
self.push_to_server_async()
self.__class__.no += 1
self.get_storage()
def record(self, data):
if data and data['input']:
......@@ -179,72 +125,22 @@ class ServerCommandRecorder(CommandRecorder, metaclass=Singleton):
data['timestamp'] = int(data['timestamp'])
self.queue.put(data)
def get_storage(self):
config = deepcopy(current_app.config["COMMAND_STORAGE"])
config['SERVICE'] = app_service
self.storage = jms_storage.get_log_storage(config)
def push_to_server_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())
logger.debug("Session command remain push: {}".format(
self.queue.qsize())
)
if not data_set:
continue
logger.debug("Send {} commands to server".format(len(data_set)))
ok = app_service.push_session_command(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("GC: Session command storage has been gc")
class ESCommandRecorder(CommandRecorder, metaclass=Singleton):
batch_size = 10
timeout = 5
no = 0
default_hosts = ["http://localhost"]
def __init__(self):
super().__init__()
self.queue = MemoryQueue()
self.stop_evt = threading.Event()
self.push_to_es_async()
self.__class__.no += 1
self.store = jms_storage.ESStore(
current_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)
ok = self.storage.bulk_save(data_set)
if not ok:
self.queue.mput(data_set)
......@@ -253,25 +149,9 @@ class ESCommandRecorder(CommandRecorder, metaclass=Singleton):
thread.start()
def session_start(self, session_id):
print("Session start: {}".format(session_id))
pass
def session_end(self, session_id):
print("Session end: {}".format(session_id))
pass
# def __del__(self):
# print("GC: ES command storage has been gc".format(self))
def get_command_recorder_class(config):
command_storage = config["COMMAND_STORAGE"]
storage_type = command_storage.get('TYPE')
if storage_type == "elasticsearch":
return ESCommandRecorder
else:
return ServerCommandRecorder
#
# def get_replay_recorder_class(config):
# ServerReplayRecorder.client = jms_storage.init(config["REPLAY_STORAGE"])
# return ServerReplayRecorder
requirements/requirements.txt
View file @ 5ec6eae5
......@@ -18,8 +18,7 @@ idna==2.6
itsdangerous==0.24
Jinja2==2.10
jmespath==0.9.3
jms-es-sdk==0.5.2
jms-storage==0.0.12
jms-storage==0.0.15
jumpserver-python-sdk==0.0.42
MarkupSafe==1.0
oss2==2.4.0
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment