#!/usr/bin/env python3 # coding: utf-8 import os import subprocess import threading import logging import logging.handlers import time import argparse import sys 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) try: from apps.jumpserver import const __version__ = const.VERSION except ImportError as e: print("Not found __version__: {}".format(e)) print("Sys path: {}".format(sys.path)) print("Python is: ") print(subprocess.call('which python', shell=True)) __version__ = 'Unknown' try: import apps print("List apps: {}".format(os.listdir('apps'))) print('apps is: {}'.format(apps)) except: pass try: from apps.jumpserver.conf import load_user_config CONFIG = load_user_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 DEBUG = CONFIG.DEBUG or False LOG_LEVEL = CONFIG.LOG_LEVEL or 'INFO' START_TIMEOUT = 40 WORKERS = 4 DAEMON = False EXIT_EVENT = threading.Event() LOCK = threading.Lock() daemon_pid_file = '' try: os.makedirs(os.path.join(BASE_DIR, "data", "static")) os.makedirs(os.path.join(BASE_DIR, "data", "media")) except: pass class LogPipe(threading.Thread): def __init__(self, name, file_path, to_stdout=False): """Setup the object with a logger and a loglevel and start the thread """ threading.Thread.__init__(self) self.daemon = False self.name = name self.file_path = file_path self.to_stdout = to_stdout self.fdRead, self.fdWrite = os.pipe() self.pipeReader = os.fdopen(self.fdRead) self.logger = self.init_logger() self.start() def init_logger(self): _logger = logging.getLogger(self.name) _logger.setLevel(logging.INFO) _formatter = logging.Formatter('%(message)s') _handler = logging.handlers.RotatingFileHandler( self.file_path, mode='a', maxBytes=5*1024*1024, backupCount=5 ) _handler.setFormatter(_formatter) _handler.setLevel(logging.INFO) _logger.addHandler(_handler) if self.to_stdout: _console = logging.StreamHandler() _console.setLevel(logging.INFO) _console.setFormatter(_formatter) _logger.addHandler(_console) return _logger def fileno(self): """Return the write file descriptor of the pipe """ return self.fdWrite def run(self): """Run the thread, logging everything. """ for line in iter(self.pipeReader.readline, ''): self.logger.info(line.strip('\n')) self.pipeReader.close() def close(self): """Close the write end of the pipe. """ os.close(self.fdWrite) def check_database_connection(): os.chdir(os.path.join(BASE_DIR, 'apps')) for i in range(60): print("Check database connection ...") code = subprocess.call("python manage.py showmigrations users ", shell=True) if code == 0: print("Database connect success") return time.sleep(1) print("Connection database failed, exist") sys.exit(10) def make_migrations(): print("Check database structure change ...") os.chdir(os.path.join(BASE_DIR, 'apps')) print("Migrate model change to database ...") subprocess.call('python3 manage.py migrate', shell=True) def collect_static(): print("Collect static files") os.chdir(os.path.join(BASE_DIR, 'apps')) command = 'python3 manage.py collectstatic --no-input -c &> /dev/null ' \ '&& echo "Collect static file done"' subprocess.call(command, shell=True) def prepare(): check_database_connection() make_migrations() collect_static() def check_pid(pid): """ Check For the existence of a unix pid. """ try: os.kill(pid, 0) except OSError: return False else: return True def get_pid_file_path(s): return os.path.join('/tmp', '{}.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): all_services = [ 'gunicorn', 'celery_ansible', 'celery_default', 'beat', 'flower' ] if s == 'all': return all_services elif s == 'gunicorn': return ['gunicorn', 'flower'] elif s == "celery": return ["celery_ansible", "celery_default"] 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() service = 'gunicorn' bind = '{}:{}'.format(HTTP_HOST, HTTP_PORT) log_format = '%(h)s %(t)s "%(r)s" %(s)s %(b)s ' pid_file = get_pid_file_path(service) cmd = [ 'gunicorn', 'jumpserver.wsgi', '-b', bind, '-k', 'gthread', '--threads', '10', '-w', str(WORKERS), '--max-requests', '4096', '--access-logformat', log_format, '-p', pid_file, '--access-logfile', '-' ] if DEBUG: cmd.append('--reload') 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', 2) 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') cmd = [ 'celery', 'worker', '-A', 'ops', '-l', 'INFO', '-c', str(num), '-Q', queue, '-n', '{}@%h'.format(queue) ] 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=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") os.environ.setdefault('PYTHONOPTIMIZE', '1') if os.getuid() == 0: os.environ.setdefault('C_FORCE_ROOT', '1') scheduler = "django_celery_beat.schedulers:DatabaseScheduler" cmd = [ 'celery', 'beat', '-A', 'ops', '-l', 'INFO', '--scheduler', scheduler, '--max-interval', '60' ] return {"cmd": cmd, 'cwd': APPS_DIR} processes = {} def watch_services(): max_retry = 3 signal.signal(signal.SIGTERM, lambda x, y: clean_up()) services_retry = defaultdict(int) stopped_services = {} def check_services(): for s, p in processes.items(): try: p.wait(timeout=1) stopped_services[s] = '' except subprocess.TimeoutExpired: stopped_services.pop(s, None) services_retry.pop(s, None) continue def retry_start_stopped_services(): for s in stopped_services: if services_retry[s] > max_retry: print("\nService start failed, exit: ", s) EXIT_EVENT.set() break print("\n> Find {} stopped, retry {}".format( s, services_retry[s] + 1) ) p = start_service(s) processes[s] = p services_retry[s] += 1 while not EXIT_EVENT.is_set(): try: with LOCK: check_services() retry_start_stopped_services() time.sleep(10) except KeyboardInterrupt: 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, "beat": get_start_beat_kwargs, "flower": get_start_flower_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') to_stdout = False if not DAEMON: to_stdout = True log_file = get_log_file_path(s) _logger = LogPipe(s, log_file, to_stdout=to_stdout) stderr = stdout = _logger kwargs.update({"stderr": stderr, "stdout": stdout}) 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): print(time.ctime()) print('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) global daemon_pid_file daemon_pid_file = get_pid_file_path('jms') context = daemon.DaemonContext( pidfile=pidfile.TimeoutPIDLockFile(daemon_pid_file), signal_map={ signal.SIGTERM: clean_up, signal.SIGHUP: 'terminate', }, ) with context: watch_services() def stop_service(s, sig=15): services_set = parse_service(s) for s in services_set: if not is_running(s): show_service_status(s) continue print("Stop service: {}".format(s)) pid = get_pid(s) os.kill(pid, sig) with LOCK: processes.pop(s, None) if s == "all": pid = get_pid('jms') os.kill(pid, sig) 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)) 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"), help="Action to run" ) parser.add_argument( "service", type=str, default="all", nargs="?", choices=("all", "gunicorn", "celery", "beat", "celery,beat", "flower"), help="The service to start", ) parser.add_argument('-d', '--daemon', nargs="?", const=1) parser.add_argument('-w', '--worker', type=int, nargs="?", const=4) 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": stop_service(srv) elif action == "restart": DAEMON = True stop_service(srv) time.sleep(5) start_services_and_watch(srv) else: show_service_status(srv)