openGauss-OM/script/killall

#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################################
# Copyright (c) 2020 Huawei Technologies Co.,Ltd.
#
# openGauss is licensed under Mulan PSL v2.
# You can use this software according to the terms
# and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS,
# WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
# ----------------------------------------------------------------------------
# Description  : killall is a program for kill process.
#############################################################################
"""

"""

import sys
import os
import pwd
import subprocess
import optparse
from gspylib.common.GaussLog import GaussLog
from base_utils.os.env_util import EnvUtil


class Closelogger():
    def __init__(self):
        pass

    def printMessage(self, msg):
        sys.stdout.write("%s\n" % msg)

    def debug(self, msg):
        pass


def init_logger():
    global logger
    LOG_DEBUG = 1
    log_path = EnvUtil.getEnv("GAUSSLOG")
    static_config_file = os.path.realpath("%s/bin/cluster_static_config")\
                         % EnvUtil.getEnv("GAUSSHOME")
    if not os.path.exists(str(log_path)) or \
            not os.path.exists(str(static_config_file)):
        logger = Closelogger()
    else:
        logFile = os.path.realpath("%s/om/%s" % (log_path, "killall_log.log"))
        logger = GaussLog(logFile, "killall", LOG_DEBUG)


def parse_command():
    """parser command"""
    parser = optparse.OptionParser(conflict_handler='resolve')
    parser.disable_interspersed_args()
    parser.add_option('-s', dest='signal', help='Sends a specified signal.')
    parser.add_option('-u', dest='user', help='Specified user')
    return parser


def exec_kill_process(user="", signal="", processname=""):
    """execute command"""
    logger.debug("start to kill process %s" % processname)
    if signal == "":
        signal = 15
    if user:
        user_uid = pwd.getpwnam("%s" % user).pw_gid
        cmd = "ps -eo uid,pid,comm | awk '{if ($1==\"%s\")print}' | awk" \
              " '{if ($3==\"%s\")print}'" % (user_uid, processname)
    else:
        cmd = "ps -eo uid,pid,comm | awk '{if ($3==\"%s\")print}'" \
              % processname
    (status, output) = subprocess.getstatusoutput(cmd)
    if status != 0:
        logger.printMessage("%s: no process found" % processname)
        sys.exit(1)
    if output:
        result_list = output.strip().split('\n')
    else:
        logger.printMessage("%s: no process found" % processname)
        sys.exit(1)
    for result in result_list:
        result = [i.strip() for i in result.strip().split(" ") if i]
        if result[2].strip() == processname:
            cmd = "kill -%s %s" % (signal, result[1])
            (status, output) = subprocess.getstatusoutput(cmd)
            if status != 0:
                cmd = "ps -eo pid | awk \'{print $1}\' | grep %s" % result[1]
                (status, output) = subprocess.getstatusoutput(cmd)
                if status == 0 and output:
                    logger.debug("Error: %s(%s): Operation not permitted"
                                 % (processname, result[1]))
                    logger.printMessage("%s(%s): Operation not permitted"
                                        % (processname, result[1]))
                else:
                    logger.debug("Error: %s: no process found" % processname)
                    logger.printMessage("%s: no process found" % processname)
                    sys.exit(1)
    logger.debug("kill process %s success" % processname)
    sys.exit(0)


def check_signal_value(signal, user, processname):
    signal_list = [i for i in range(1, 65)]
    if int(signal) not in signal_list:
        if user:
            user_uid = pwd.getpwnam("%s" % user).pw_gid
            cmd = "ps -eo uid,pid,comm | awk '{if ($1==\"%s\")print}' | awk " \
                  "'{if ($3==\"%s\")print}'" % (user_uid, processname)
        else:
            cmd = "ps -eo uid,pid,comm | awk '{if ($3==\"%s\")print}'"\
                  % processname
        (status, output) = subprocess.getstatusoutput(cmd)
        # The query result does not exist, and the status is 1
        if status != 0:
            logger.printMessage("%s: no process found" % processname)
            sys.exit(1)
        if output:
            result_list = output.strip().split('\n')
        else:
            logger.printMessage("%s: no process found" % processname)
            sys.exit(1)
        for result in result_list:
            result = [i.strip() for i in result.strip().split(" ") if i]
            if result[2].strip() == processname:
                logger.debug("Error: %s(%s): Invalid argument"
                             % (processname, result[1]))
                logger.printMessage("%s(%s): Invalid argument"
                                    % (processname, result[1]))
        logger.printMessage("%s: no process found" % processname)
        sys.exit(1)
    else:
        return int(signal)


if __name__ == '__main__':
    """main"""
    init_logger()
    signal_type = ""
    for value in sys.argv[1:]:
        if value.strip().startswith("-") and value.count("-") == 1 and \
                value.strip().split("-")[1].isdigit():
            signal_type = value.strip().split("-")[1]
            sys.argv.remove(value)
            break
    parser = parse_command()
    opts, args = parser.parse_args()
    signal = ""
    user = ""
    process = ""
    if opts.user:
        user = opts.user.strip()
    if args:
        process = args[0]
    if opts.signal:
        signal = opts.signal.strip()
    elif signal_type:
        signal_type = check_signal_value(signal_type, user, process)
        signal = signal_type
    exec_kill_process(user=user, signal=signal, processname=process)