mirror of https://github.com/jumpserver/jumpserver
583 lines
16 KiB
Python
Executable File
583 lines
16 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
# coding: utf-8
|
||
|
||
import os
|
||
import subprocess
|
||
import threading
|
||
import datetime
|
||
import logging
|
||
import logging.handlers
|
||
import psutil
|
||
import time
|
||
import argparse
|
||
import sys
|
||
import shutil
|
||
import signal
|
||
from collections import defaultdict
|
||
import daemon
|
||
from daemon import pidfile
|
||
|
||
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
||
sys.path.insert(0, BASE_DIR)
|
||
|
||
logging.basicConfig(level=logging.DEBUG, format="%(asctime)s %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
|
||
|
||
try:
|
||
from apps.jumpserver import const
|
||
__version__ = const.VERSION
|
||
except ImportError as e:
|
||
print("Not found __version__: {}".format(e))
|
||
print("Python is: ")
|
||
logging.info(subprocess.call('which python', shell=True))
|
||
__version__ = 'Unknown'
|
||
sys.exit(1)
|
||
|
||
try:
|
||
from apps.jumpserver.const import CONFIG
|
||
except ImportError as e:
|
||
print("Import error: {}".format(e))
|
||
print("Could not find config file, `cp config_example.yml config.yml`")
|
||
sys.exit(1)
|
||
|
||
os.environ["PYTHONIOENCODING"] = "UTF-8"
|
||
APPS_DIR = os.path.join(BASE_DIR, 'apps')
|
||
LOG_DIR = os.path.join(BASE_DIR, 'logs')
|
||
TMP_DIR = os.path.join(BASE_DIR, 'tmp')
|
||
HTTP_HOST = CONFIG.HTTP_BIND_HOST or '127.0.0.1'
|
||
HTTP_PORT = CONFIG.HTTP_LISTEN_PORT or 8080
|
||
WS_PORT = CONFIG.WS_LISTEN_PORT or 8082
|
||
DEBUG = CONFIG.DEBUG or False
|
||
LOG_LEVEL = CONFIG.LOG_LEVEL or 'INFO'
|
||
|
||
START_TIMEOUT = 40
|
||
WORKERS = 4
|
||
DAEMON = False
|
||
LOG_KEEP_DAYS = 7
|
||
logging.basicConfig(
|
||
format='%(asctime)s %(message)s', level=logging.INFO,
|
||
datefmt='%Y-%m-%d %H:%M:%S'
|
||
)
|
||
|
||
EXIT_EVENT = threading.Event()
|
||
LOCK = threading.Lock()
|
||
files_preserve = []
|
||
STOP_TIMEOUT = 10
|
||
|
||
logger = logging.getLogger()
|
||
|
||
try:
|
||
os.makedirs(os.path.join(BASE_DIR, "data", "static"))
|
||
os.makedirs(os.path.join(BASE_DIR, "data", "media"))
|
||
except:
|
||
pass
|
||
|
||
|
||
def check_database_connection():
|
||
os.chdir(os.path.join(BASE_DIR, 'apps'))
|
||
for i in range(60):
|
||
logging.info("Check database connection ...")
|
||
code = subprocess.call("python manage.py showmigrations users ", shell=True)
|
||
if code == 0:
|
||
logging.info("Database connect success")
|
||
return
|
||
time.sleep(1)
|
||
logging.error("Connection database failed, exist")
|
||
sys.exit(10)
|
||
|
||
|
||
def check_migrations():
|
||
apps_dir = os.path.join(BASE_DIR, 'apps')
|
||
code = subprocess.call("python manage.py showmigrations | grep '\[.\]' | grep -v '\[X\]'", shell=True, cwd=apps_dir)
|
||
|
||
if code == 1:
|
||
return
|
||
# for i in range(3):
|
||
# print("!!! Warning: Has SQL migrations not perform, 有 SQL 变更没有执行")
|
||
# print("You should run `./PROC upgrade_db` first, 请先运行 ./PROC upgrade_db, 进行表结构变更")
|
||
# sys.exit(1)
|
||
|
||
|
||
def perform_db_migrate():
|
||
logging.info("Check database structure change ...")
|
||
os.chdir(os.path.join(BASE_DIR, 'apps'))
|
||
logging.info("Migrate model change to database ...")
|
||
subprocess.call('python3 manage.py migrate', shell=True)
|
||
|
||
|
||
def collect_static():
|
||
logging.info("Collect static files")
|
||
os.chdir(os.path.join(BASE_DIR, 'apps'))
|
||
command = 'python3 manage.py collectstatic --no-input -c &> /dev/null '
|
||
subprocess.call(command, shell=True)
|
||
logging.info("Collect static files done")
|
||
|
||
|
||
def prepare():
|
||
check_database_connection()
|
||
check_migrations()
|
||
upgrade_db()
|
||
|
||
|
||
def check_pid(pid):
|
||
""" Check For the existence of a unix pid. """
|
||
try:
|
||
os.kill(pid, 0)
|
||
except (OSError, ProcessLookupError):
|
||
return False
|
||
else:
|
||
return True
|
||
|
||
|
||
def get_pid_file_path(s):
|
||
return os.path.join(TMP_DIR, '{}.pid'.format(s))
|
||
|
||
|
||
def get_log_file_path(s):
|
||
return os.path.join(LOG_DIR, '{}.log'.format(s))
|
||
|
||
|
||
def get_pid_from_file(path):
|
||
if os.path.isfile(path):
|
||
with open(path) as f:
|
||
try:
|
||
return int(f.read().strip())
|
||
except ValueError:
|
||
return 0
|
||
return 0
|
||
|
||
|
||
def get_pid(s):
|
||
pid_file = get_pid_file_path(s)
|
||
return get_pid_from_file(pid_file)
|
||
|
||
|
||
def is_running(s, unlink=True):
|
||
pid_file = get_pid_file_path(s)
|
||
|
||
if os.path.isfile(pid_file):
|
||
pid = get_pid(s)
|
||
if pid == 0:
|
||
return False
|
||
elif check_pid(pid):
|
||
return True
|
||
|
||
if unlink:
|
||
os.unlink(pid_file)
|
||
return False
|
||
|
||
|
||
def parse_service(s):
|
||
web_services = ['gunicorn', 'flower', 'daphne']
|
||
celery_services = [
|
||
"celery_ansible", "celery_default", "celery_node_tree",
|
||
"celery_check_asset_perm_expired", "celery_heavy_tasks"
|
||
]
|
||
task_services = celery_services + ['beat']
|
||
all_services = web_services + task_services
|
||
if s == 'all':
|
||
return all_services
|
||
elif s == "web":
|
||
return web_services
|
||
elif s == 'ws':
|
||
return ['daphne']
|
||
elif s == "task":
|
||
return task_services
|
||
elif s == "celery":
|
||
return celery_services
|
||
elif "," in s:
|
||
services = set()
|
||
for i in s.split(','):
|
||
services.update(parse_service(i))
|
||
return services
|
||
else:
|
||
return [s]
|
||
|
||
|
||
def get_start_gunicorn_kwargs():
|
||
print("\n- Start Gunicorn WSGI HTTP Server")
|
||
prepare()
|
||
bind = '{}:{}'.format(HTTP_HOST, HTTP_PORT)
|
||
log_format = '%(h)s %(t)s %(L)ss "%(r)s" %(s)s %(b)s '
|
||
|
||
cmd = [
|
||
'gunicorn', 'jumpserver.wsgi',
|
||
'-b', bind,
|
||
'-k', 'gthread',
|
||
'--threads', '10',
|
||
'-w', str(WORKERS),
|
||
'--max-requests', '4096',
|
||
'--access-logformat', log_format,
|
||
'--access-logfile', '-'
|
||
]
|
||
|
||
if DEBUG:
|
||
cmd.append('--reload')
|
||
return {'cmd': cmd, 'cwd': APPS_DIR}
|
||
|
||
|
||
def get_start_daphne_kwargs():
|
||
print("\n- Start Daphne ASGI WS Server")
|
||
cmd = [
|
||
'daphne', 'jumpserver.asgi:application',
|
||
'-b', HTTP_HOST,
|
||
'-p', str(WS_PORT),
|
||
]
|
||
return {'cmd': cmd, 'cwd': APPS_DIR}
|
||
|
||
|
||
def get_start_celery_ansible_kwargs():
|
||
print("\n- Start Celery as Distributed Task Queue: Ansible")
|
||
return get_start_worker_kwargs('ansible', 4)
|
||
|
||
|
||
def get_start_celery_default_kwargs():
|
||
print("\n- Start Celery as Distributed Task Queue: Celery")
|
||
return get_start_worker_kwargs('celery', 4)
|
||
|
||
|
||
def get_start_celery_node_tree_kwargs():
|
||
print("\n- Start Celery as Distributed Task Queue: NodeTree")
|
||
return get_start_worker_kwargs('node_tree', 2)
|
||
|
||
|
||
def get_start_celery_heavy_tasks_kwargs():
|
||
print("\n- Start Celery as Distributed Task Queue: HeavyTasks")
|
||
return get_start_worker_kwargs('celery_heavy_tasks', 1)
|
||
|
||
|
||
def get_start_celery_check_asset_perm_expired_kwargs():
|
||
print("\n- Start Celery as Distributed Task Queue: CheckAseetPermissionExpired")
|
||
return get_start_worker_kwargs('celery_check_asset_perm_expired', 1)
|
||
|
||
|
||
def get_start_worker_kwargs(queue, num):
|
||
# Todo: Must set this environment, otherwise not no ansible result return
|
||
os.environ.setdefault('PYTHONOPTIMIZE', '1')
|
||
os.environ.setdefault('ANSIBLE_FORCE_COLOR', 'True')
|
||
|
||
if os.getuid() == 0:
|
||
os.environ.setdefault('C_FORCE_ROOT', '1')
|
||
server_hostname = os.environ.get("SERVER_HOSTNAME")
|
||
if not server_hostname:
|
||
server_hostname = '%h'
|
||
|
||
cmd = [
|
||
'celery', 'worker',
|
||
'-A', 'ops',
|
||
'-l', 'INFO',
|
||
'-c', str(num),
|
||
'-Q', queue,
|
||
'-n', '{}@{}'.format(queue, server_hostname)
|
||
]
|
||
return {"cmd": cmd, "cwd": APPS_DIR}
|
||
|
||
|
||
def get_start_flower_kwargs():
|
||
print("\n- Start Flower as Task Monitor")
|
||
if os.getuid() == 0:
|
||
os.environ.setdefault('C_FORCE_ROOT', '1')
|
||
|
||
cmd = [
|
||
'celery', 'flower',
|
||
'-A', 'ops',
|
||
'-l', 'INFO',
|
||
'--url_prefix=/core/flower',
|
||
'--auto_refresh=False',
|
||
'--max_tasks=1000',
|
||
'--tasks_columns=uuid,name,args,state,received,started,runtime,worker'
|
||
]
|
||
return {"cmd": cmd, "cwd": APPS_DIR}
|
||
|
||
|
||
def get_start_beat_kwargs():
|
||
print("\n- Start Beat as Periodic Task Scheduler")
|
||
utils_dir = os.path.join(BASE_DIR, 'utils')
|
||
cmd = [
|
||
sys.executable, 'start_celery_beat.py',
|
||
]
|
||
return {"cmd": cmd, 'cwd': utils_dir}
|
||
|
||
|
||
processes = {}
|
||
|
||
|
||
def watch_services():
|
||
max_retry = 3
|
||
services_retry = defaultdict(int)
|
||
stopped_services = {}
|
||
|
||
def check_services():
|
||
now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
||
for s, p in processes.items():
|
||
print("{} Check service status: {} -> ".format(now, s), end='')
|
||
try:
|
||
p.wait(timeout=1) # 不wait,子进程可能无法回收
|
||
except subprocess.TimeoutExpired:
|
||
pass
|
||
ok = is_running(s)
|
||
if not ok:
|
||
stopped_services[s] = ''
|
||
print("stopped with code: {}({})".format(p.returncode, p.pid))
|
||
else:
|
||
print("running at {}".format(p.pid))
|
||
stopped_services.pop(s, None)
|
||
services_retry.pop(s, None)
|
||
|
||
def retry_start_stopped_services():
|
||
for s in stopped_services:
|
||
if services_retry[s] > max_retry:
|
||
logging.info("Service start failed, exit: ", s)
|
||
EXIT_EVENT.set()
|
||
break
|
||
|
||
p = start_service(s)
|
||
logging.info("> Find {} stopped, retry {}, {}".format(
|
||
s, services_retry[s] + 1, p.pid)
|
||
)
|
||
processes[s] = p
|
||
services_retry[s] += 1
|
||
|
||
def rotate_log_if_need():
|
||
now = datetime.datetime.now()
|
||
tm = now.strftime('%H:%M')
|
||
if tm != '23:59':
|
||
return
|
||
suffix = now.strftime('%Y-%m-%d')
|
||
services = list(processes.keys())
|
||
services.append('jms')
|
||
|
||
for s in services:
|
||
log_path = get_log_file_path(s)
|
||
log_dir = os.path.dirname(log_path)
|
||
filename = os.path.basename(log_path)
|
||
pre_log_dir = os.path.join(log_dir, suffix)
|
||
if not os.path.exists(pre_log_dir):
|
||
os.mkdir(pre_log_dir)
|
||
|
||
pre_log_path = os.path.join(pre_log_dir, filename)
|
||
if os.path.isfile(log_path) and not os.path.isfile(pre_log_path):
|
||
logging.info("Rotate log file: {} => {}".format(log_path, pre_log_path))
|
||
shutil.copy(log_path, pre_log_path)
|
||
with open(log_path, 'w') as f:
|
||
pass
|
||
some_days_ago = now - datetime.timedelta(days=LOG_KEEP_DAYS)
|
||
days_ago_dir = os.path.join(LOG_DIR, some_days_ago.strftime('%Y-%m-%d'))
|
||
if os.path.exists(days_ago_dir):
|
||
logger.info("Remove old log: {}".format(days_ago_dir))
|
||
shutil.rmtree(days_ago_dir, ignore_errors=True)
|
||
|
||
while not EXIT_EVENT.is_set():
|
||
try:
|
||
with LOCK:
|
||
check_services()
|
||
retry_start_stopped_services()
|
||
rotate_log_if_need()
|
||
time.sleep(30)
|
||
except KeyboardInterrupt:
|
||
print("Start stop service")
|
||
time.sleep(1)
|
||
break
|
||
clean_up()
|
||
|
||
|
||
def start_service(s):
|
||
services_kwargs = {
|
||
"gunicorn": get_start_gunicorn_kwargs,
|
||
"celery_ansible": get_start_celery_ansible_kwargs,
|
||
"celery_default": get_start_celery_default_kwargs,
|
||
"celery_node_tree": get_start_celery_node_tree_kwargs,
|
||
"celery_heavy_tasks": get_start_celery_heavy_tasks_kwargs,
|
||
"celery_check_asset_perm_expired": get_start_celery_check_asset_perm_expired_kwargs,
|
||
"beat": get_start_beat_kwargs,
|
||
"flower": get_start_flower_kwargs,
|
||
"daphne": get_start_daphne_kwargs,
|
||
}
|
||
|
||
kwargs = services_kwargs.get(s)()
|
||
pid_file = get_pid_file_path(s)
|
||
|
||
if os.path.isfile(pid_file):
|
||
os.unlink(pid_file)
|
||
cmd = kwargs.pop('cmd')
|
||
|
||
log_file_path = get_log_file_path(s)
|
||
log_file_f = open(log_file_path, 'a')
|
||
files_preserve.append(log_file_f)
|
||
kwargs['stderr'] = log_file_f
|
||
kwargs['stdout'] = log_file_f
|
||
p = subprocess.Popen(cmd, **kwargs)
|
||
with open(pid_file, 'w') as f:
|
||
f.write(str(p.pid))
|
||
return p
|
||
|
||
|
||
def start_services_and_watch(s):
|
||
logging.info(time.ctime())
|
||
logging.info('JumpServer version {}, more see https://www.jumpserver.org'.format(
|
||
__version__)
|
||
)
|
||
|
||
services_set = parse_service(s)
|
||
for i in services_set:
|
||
if is_running(i):
|
||
show_service_status(i)
|
||
continue
|
||
p = start_service(i)
|
||
time.sleep(2)
|
||
processes[i] = p
|
||
|
||
if not DAEMON:
|
||
watch_services()
|
||
else:
|
||
show_service_status(s)
|
||
context = get_daemon_context()
|
||
with context:
|
||
watch_services()
|
||
|
||
|
||
def get_daemon_context():
|
||
daemon_pid_file = get_pid_file_path('jms')
|
||
daemon_log_f = open(get_log_file_path('jms'), 'a')
|
||
files_preserve.append(daemon_log_f)
|
||
context = daemon.DaemonContext(
|
||
pidfile=pidfile.TimeoutPIDLockFile(daemon_pid_file),
|
||
signal_map={
|
||
signal.SIGTERM: lambda x, y: clean_up(),
|
||
signal.SIGHUP: 'terminate',
|
||
},
|
||
stdout=daemon_log_f,
|
||
stderr=daemon_log_f,
|
||
files_preserve=files_preserve,
|
||
detach_process=True,
|
||
)
|
||
return context
|
||
|
||
|
||
def stop_service(srv, sig=15):
|
||
services_set = parse_service(srv)
|
||
for s in services_set:
|
||
if not is_running(s):
|
||
show_service_status(s)
|
||
continue
|
||
print("Stop service: {}".format(s), end='')
|
||
pid = get_pid(s)
|
||
os.kill(pid, sig)
|
||
with LOCK:
|
||
process = processes.pop(s, None)
|
||
if process is None:
|
||
try:
|
||
process = psutil.Process(pid)
|
||
except:
|
||
pass
|
||
if process is None:
|
||
print("\033[31m No process found\033[0m")
|
||
continue
|
||
try:
|
||
process.wait(1)
|
||
except:
|
||
pass
|
||
for i in range(STOP_TIMEOUT):
|
||
if i == STOP_TIMEOUT - 1:
|
||
print("\033[31m Error\033[0m")
|
||
if not is_running(s):
|
||
print("\033[32m Ok\033[0m")
|
||
break
|
||
else:
|
||
time.sleep(1)
|
||
continue
|
||
|
||
if srv == "all":
|
||
stop_daemon_service()
|
||
|
||
|
||
def stop_daemon_service():
|
||
pid = get_pid('jms')
|
||
if pid and check_pid(pid):
|
||
os.kill(pid, 15)
|
||
|
||
|
||
def stop_multi_services(services):
|
||
for s in services:
|
||
stop_service(s, sig=9)
|
||
|
||
|
||
def stop_service_force(s):
|
||
stop_service(s, sig=9)
|
||
|
||
|
||
def clean_up():
|
||
if not EXIT_EVENT.is_set():
|
||
EXIT_EVENT.set()
|
||
processes_dump = {k: v for k, v in processes.items()}
|
||
for s1, p1 in processes_dump.items():
|
||
stop_service(s1)
|
||
p1.wait()
|
||
|
||
|
||
def show_service_status(s):
|
||
services_set = parse_service(s)
|
||
for ns in services_set:
|
||
if is_running(ns):
|
||
pid = get_pid(ns)
|
||
print("{} is running: {}".format(ns, pid))
|
||
else:
|
||
print("{} is stopped".format(ns))
|
||
|
||
|
||
def upgrade_db():
|
||
collect_static()
|
||
perform_db_migrate()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
parser = argparse.ArgumentParser(
|
||
description="""
|
||
Jumpserver service control tools;
|
||
|
||
Example: \r\n
|
||
|
||
%(prog)s start all -d;
|
||
"""
|
||
)
|
||
parser.add_argument(
|
||
'action', type=str,
|
||
choices=("start", "stop", "restart", "status", "upgrade_db"),
|
||
help="Action to run"
|
||
)
|
||
parser.add_argument(
|
||
"service", type=str, default="all", nargs="?",
|
||
choices=("all", "web", "task", "gunicorn", "celery", "beat", "celery,beat", "flower", "ws"),
|
||
help="The service to start",
|
||
)
|
||
parser.add_argument('-d', '--daemon', nargs="?", const=1)
|
||
parser.add_argument('-w', '--worker', type=int, nargs="?", const=4)
|
||
parser.add_argument('-f', '--force', nargs="?", const=1)
|
||
args = parser.parse_args()
|
||
if args.daemon:
|
||
DAEMON = True
|
||
|
||
if args.worker:
|
||
WORKERS = args.worker
|
||
|
||
action = args.action
|
||
srv = args.service
|
||
|
||
if action == "start":
|
||
start_services_and_watch(srv)
|
||
os._exit(0)
|
||
elif action == "stop":
|
||
print("Stop service")
|
||
if args.force:
|
||
stop_service_force(srv)
|
||
else:
|
||
stop_service(srv)
|
||
elif action == "restart":
|
||
DAEMON = True
|
||
stop_service(srv)
|
||
time.sleep(5)
|
||
start_services_and_watch(srv)
|
||
elif action == "upgrade_db":
|
||
upgrade_db()
|
||
else:
|
||
show_service_status(srv)
|