jms

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

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

import django
from django.core import management
from django.db.utils import OperationalError

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
APP_DIR = os.path.join(BASE_DIR, 'apps')

os.chdir(APP_DIR)
sys.path.insert(0, APP_DIR)
os.environ.setdefault("DJANGO_SETTINGS_MODULE", "jumpserver.settings")
django.setup()

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

try:
    from 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 jumpserver.const import CONFIG
    from common.utils.file import download_file
    from common.utils import make_dirs
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:
    make_dirs(os.path.join(BASE_DIR, "data", "static"))
    make_dirs(os.path.join(BASE_DIR, "data", "media"))
except:
    pass


def check_database_connection():
    for i in range(60):
        logging.info(f"Check database connection: {i}")
        try:
            management.call_command('check', '--database', 'default')
            logging.info("Database connect success")
            return
        except OperationalError:
            logging.info('Database not setup, retry')
        except Exception as exc:
            logging.error('Unexpect error occur: {}'.format(str(exc)))
        time.sleep(1)
    logging.error("Connection database failed, exit")
    sys.exit(10)


def expire_caches():
    try:
        management.call_command('expire_caches')
    except:
        pass


def perform_db_migrate():
    logging.info("Check database structure change ...")
    logging.info("Migrate model change to database ...")
    try:
        management.call_command('migrate')
    except Exception:
        logging.error('Perform migrate failed, exit', exc_info=True)
        sys.exit(11)


def collect_static():
    logging.info("Collect static files")
    try:
        management.call_command('collectstatic', '--no-input', '-c', verbosity=0, interactive=False)
        logging.info("Collect static files done")
    except:
        pass


def compile_i18n_file():
    django_mo_file = os.path.join(BASE_DIR, 'apps', 'locale', 'zh', 'LC_MESSAGES', 'django.mo')
    if os.path.exists(django_mo_file):
        return
    os.chdir(os.path.join(BASE_DIR, 'apps'))
    management.call_command('compilemessages', verbosity=0, interactive=False)
    logging.info("Compile i18n files done")


def download_ip_db():
    db_base_dir = os.path.join(APP_DIR, 'common', 'utils', 'ip')
    db_path_url_mapper = {
        ('geoip', 'GeoLite2-City.mmdb'): 'https://jms-pkg.oss-cn-beijing.aliyuncs.com/ip/GeoLite2-City.mmdb',
        ('ipip', 'ipipfree.ipdb'): 'https://jms-pkg.oss-cn-beijing.aliyuncs.com/ip/ipipfree.ipdb'
    }
    for p, src in db_path_url_mapper.items():
        path = os.path.join(db_base_dir, *p)
        if os.path.isfile(path) and os.path.getsize(path) > 1000:
            continue
        print("Download ip db: {}".format(path))
        download_file(src, path)


def install_builtin_applets():
    logging.info("Install builtin applets")
    try:
        management.call_command('install_builtin_applets', verbosity=0)
    except Exception as e:
        logging.error("Install builtin applets err: {}".format(e))


def upgrade_db():
    collect_static()
    perform_db_migrate()


def prepare():
    check_database_connection()
    upgrade_db()
    expire_caches()
    download_ip_db()
    install_builtin_applets()


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

    start_args = []
    if args.daemon:
        start_args.append('--daemon')
    if args.force:
        start_args.append('--force')
    if args.worker:
        start_args.extend(['--worker', str(args.worker)])
    else:
        worker = os.environ.get('CORE_WORKER')
        if isinstance(worker, str) and worker.isdigit():
            start_args.extend(['--worker', worker])

    try:
        management.call_command(action, *services, *start_args)
    except KeyboardInterrupt:
        logging.info('Cancel ...')
        time.sleep(2)
    except Exception as exc:
        logging.error("Start service error {}: {}".format(services, exc))
        time.sleep(2)


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", "collect_static"),
        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="?")
    parser.add_argument('-f', '--force', nargs="?", const=True)

    args = parser.parse_args()

    action = args.action
    if action == "upgrade_db":
        upgrade_db()
    elif action == "collect_static":
        collect_static()
    else:
        start_services()