hbyd_ueba/install.py

#!/usr/bin/env python
# -*-coding:utf-8 -*-

import sys
import os
import json
import shutil
import traceback
import logging
from appManager.module.global_api import ComponentHelper
from appsUtils.confutil import ConfUtil


BASE_PATH = os.path.split(os.path.realpath(__file__))[0]
DSTPATH = sys.argv[1]
ISUPDATE = sys.argv[2]
CUR_PATH = os.path.normpath(os.path.dirname(os.path.abspath(__file__)))
APP_NAME = 'uebaMetricsAnalysis'

def get_logger(logger_name=APP_NAME, logger_level=logging.INFO):
    """日志"""
    import logging.handlers
    logger = logging.getLogger(logger_name)
    formatter = logging.Formatter(
        '%(asctime)s %(levelname)s %(filename)s %(funcName)s:%(lineno)d %(message)s')
    logger.setLevel(logger_level)
    logger_file = os.path.normpath(os.path.join('/home/master', 'logs'))
    if not os.path.exists(logger_file):
        os.mkdir(logger_file)
    if not logger.handlers:
        file_handler = logging.handlers.TimedRotatingFileHandler(
            os.path.normpath(logger_file + os.sep + logger_name + '.log'), 'midnight')
        file_handler.suffix = "%Y-%m-%d"
        file_handler.setLevel(logger_level)
        file_handler.setFormatter(formatter)
        logger.addHandler(file_handler)
    return logger


logger = get_logger("uebaMetricsAnalysis")


def installDBSchema(exec_sql):
    try:
        confutil = ConfUtil()
        curPgConf = confutil.getPostgresqlConf()
        sqlPath = os.path.join(CUR_PATH, APP_NAME, 'sql')

        for sqlfile in exec_sql:
            if sqlfile == "pg_update.sql":
                continue
            sqlPath = os.path.join(CUR_PATH, APP_NAME, 'sql', sqlfile)
            cmd = 'psql -U %s %s<"%s"' % (curPgConf["username"], curPgConf["database"], sqlPath,)
            logger.info(cmd)
            r = os.popen(cmd)
            logger.info(sqlfile + '   ' + str(r.readlines()))
        logger.info('create schema successfully.')
    except Exception as e:
        logger.info(str(e))

def add_task():
    # 组件批量注册任务
    ch = ComponentHelper(APP_NAME)
    logger.info("开始注册定时任务")
    ch_ret = ch.add_task()
    if ch_ret['status'] != 200:
        logger.info('add task failed, reason is %s' % ch_ret['msg'])
    else:
        logger.info('add task successfully.')
    logger.info("定时任务注册成功")


class Install():
    def __init__(self):
        self.app_name = self.__get_app_name()

    def install(self):
        try:
            installDBSchema(["pg_struct.sql"])
            add_task()
            logger.info('>>>安装结束！！！')
        except Exception as e:
            logger.error("install failure! reason: %s" % traceback.format_exc())
            self.install_rollback()
            # 此处写自定义的状态码，如100
            code = 100
            sys.exit(code)

    def upgrade(self):
        try:
            # 第一步，拷贝目录
            shutil.rmtree(DSTPATH)
            self.copy_files()
            # 第二步，写自己升级的逻辑
        except Exception as e:
            logger.info(traceback.format_exc())
            self.upgrade_rollback()
            sys.exit(1)

    def copy_files(self):
        src_path = os.path.join(CUR_PATH, self.app_name)
        shutil.copytree(src_path, DSTPATH)

    def install_rollback(self):
        """安装回滚操作"""
        try:
            logger.error("rollback app folder")
        except Exception as e:
            logger.error("roll back failed")
            logger.error(traceback.format_exc())

    def upgrade_rollback(self):
        """升级回滚操作"""
        try:
            logger.error("rollback app folder")
        except Exception as e:
            logger.error("roll back failed")
            logger.error(traceback.format_exc())

    def __get_app_name(self):
        with open(os.path.join(CUR_PATH, list(os.walk(CUR_PATH))[0][1][0], "package.json"), "r+") as f:
            pkgConfig = json.load(f)
        return pkgConfig["name"]


def main():
    try:
        flag = sys.argv[2]
        install_instance = Install()
        if flag == '0':
            logger.info("=" * 20 + "开始安装" + "=" * 20)
            install_instance.install()
        if flag == '1':
            logger.info("=" * 20 + "开始升级" + "=" * 20)
            install_instance.upgrade()
        logger.info("main 执行完成")
        sys.exit(0)
    except Exception as e:
        logger.info("install failure, reason: %s" % traceback.format_exc())
        raise e


if __name__ == '__main__':
    main()