jumpserver/jms

#!/usr/bin/env python3
# coding: utf-8

import os
import subprocess
import logging
import logging.handlers
import time
import argparse
import sys

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(sys.executable)
    __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"

logging.basicConfig(
    format='%(asctime)s %(message)s', level=logging.INFO,
    datefmt='%Y-%m-%d %H:%M:%S'
)

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, exit")
    sys.exit(10)


def check_migrations():
    _apps_dir = os.path.join(BASE_DIR, 'apps')
    _cmd = "python manage.py showmigrations | grep '\[.\]' | grep -v '\[X\]'"
    _code = subprocess.call(_cmd, 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 expire_caches():
    _apps_dir = os.path.join(BASE_DIR, 'apps')
    _code = subprocess.call("python manage.py expire_caches", shell=True, cwd=_apps_dir)

    if _code == 1:
        return


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 ...")
    _code = subprocess.call('python3 manage.py migrate', shell=True)
    if _code == 0:
        return
    logging.error('Perform migrate failed, exit')
    sys.exit(11)


def collect_static():
    logging.info("Collect static files")
    os.chdir(os.path.join(BASE_DIR, 'apps'))
    _cmd = 'python3 manage.py collectstatic --no-input -c &> /dev/null '
    subprocess.call(_cmd, shell=True)
    logging.info("Collect static files done")


def upgrade_db():
    collect_static()
    perform_db_migrate()


def prepare():
    # installer(check) & k8s(no check)
    check_database_connection()
    check_migrations()
    upgrade_db()
    expire_caches()


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(
        "services", type=str, default='all', nargs="*",
        choices=("all", "web", "task"),
        help="The service to start",
    )
    parser.add_argument('-d', '--daemon', nargs="?", const=True)
    parser.add_argument('-w', '--worker', type=int, nargs="?", default=4)
    parser.add_argument('-f', '--force', nargs="?", const=True)

    args = parser.parse_args()

    action = args.action
    if action == "upgrade_db":
        upgrade_db()
    else:
        services = args.services if isinstance(args.services, list) else [args.services]
        if action == 'start' and {'all', 'web'} & set(services):
            prepare()

        services_string = ' '.join(services)
        cmd = f'python manage.py {args.action} {services_string}'
        if args.daemon:
            cmd += ' --daemon'
        if args.worker:
            cmd += f' --worker {args.worker}'
        if args.force:
            cmd += ' --force'
        apps_dir = os.path.join(BASE_DIR, 'apps')

        try:
            # processes: main(3s) -> call(0.25s) -> service -> sub-process
            code = subprocess.run(cmd, shell=True, cwd=apps_dir)
        except KeyboardInterrupt:
            time.sleep(1)
            pass