jumpserver/receptor

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

import argparse
import logging
import shutil
import subprocess
import os
import signal
import tempfile

import psutil
from psutil import NoSuchProcess

ANSIBLE_RUNNER_COMMAND = "ansible-runner"
DEFAULT_SHARE_DIR = "data/share"
DEFAULT_CONTROL_SOCK_PATH = os.path.join(DEFAULT_SHARE_DIR, "control.sock")

logger = logging.getLogger(__name__)


class ReceptorService:
    def __init__(self):
        self.pid_file = "tmp/receptor.pid"
        self.receptor_command = [
            'receptor',
            '--local-only',
            '--log-level',
            'level=Debug',
            '--node', 'id=primary',
            '--control-service',
            'service=control',
            'filename={}'.format(DEFAULT_CONTROL_SOCK_PATH),
            '--work-command',
            'worktype={}'.format(ANSIBLE_RUNNER_COMMAND),
            'command={}'.format(ANSIBLE_RUNNER_COMMAND),
            'params=worker',
            'allowruntimeparams=true',
            '--work-command',
            'worktype={}'.format("kill"),
            'command={}'.format("python"),
            "params=receptor kill",
            'allowruntimeparams=true'
        ]

    @staticmethod
    def before_start():
        os.makedirs(os.path.join(DEFAULT_SHARE_DIR), exist_ok=True)
        status_dir = os.path.join(tempfile.gettempdir(), "receptor")
        if os.path.exists(status_dir):
            shutil.rmtree(status_dir)

    def start(self):
        self.before_start()
        if os.path.exists(self.pid_file):
            with open(self.pid_file, 'r') as f:
                pid_str = f.read()
            try:
                pid = int(pid_str)
                os.kill(pid, 0)
                print("\n- Receptor service is already running.")
                return
            except ProcessLookupError:
                print("\n- PID file exists but process does not, starting Receptor...")
            except ValueError:
                print("\n- PID file is corrupted, starting Receptor...")
                os.remove(self.pid_file)

        os.environ.update({'LOCAL_CONNECTION_ENABLED': '1'})
        process = subprocess.Popen(self.receptor_command)
        with open(self.pid_file, 'w') as f:
            f.write(str(process.pid))
        print("\n- Receptor service started successfully.")

        def exit_handler(signum, frame):
            process.terminate()
            process.kill()

        signal.signal(signal.SIGINT, exit_handler)
        signal.signal(signal.SIGTERM, exit_handler)
        process.wait()

    def stop(self):
        if not os.path.exists(self.pid_file):
            print("\n- Receptor service is not running.")
            return
        with open(self.pid_file, 'r') as f:
            pid = int(f.read())
        try:
            os.kill(pid, signal.SIGTERM)
            os.remove(self.pid_file)
            print("\n- Receptor service stopped successfully.")
        except ProcessLookupError:
            print("\n- Failed to stop Receptor service: Process does not exist.")
            os.remove(self.pid_file)

    def restart(self):
        self.stop()
        self.start()

    def status(self):
        if os.path.exists(self.pid_file):
            with open(self.pid_file, 'r') as f:
                pid_str = f.read()
            try:
                pid = int(pid_str)
                os.kill(pid, 0)
                print("\n- Receptor service is running.")
                return
            except ProcessLookupError:
                print("\n- Receptor service is not running.")
        else:
            print("\n- Receptor service is not running.")


def handle_receptor_action(args):
    action = args.action
    srv = ReceptorService()
    if action == "start":
        srv.start()
    elif action == 'stop':
        srv.stop()
    elif action == "restart":
        srv.restart()
    elif action == "status":
        srv.status()
    elif action == "kill":
        kill_progress_tree()


def kill_progress_tree(pid=None):
    if not pid:
        try:
            pid_input = input()
            pid = int(pid_input)
        except Exception as e:
            logger.error(e)
            return

    logger.info("progress {} will be kill".format(pid))

    try:
        current_process = psutil.Process(pid)
    except NoSuchProcess as e:
        logger.error(e)
        return

    children = current_process.children(recursive=True)
    for child in children:
        if child.pid == 1:
            continue
        if child.name() != 'ssh':
            continue
        try:
            child.kill()
        except Exception as e:
            logger.error(e)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="""
           Jumpserver receptor service control tools;
           """
    )
    parser.add_argument(
        'action', type=str,
        choices=("start", "stop", "restart", "status", "kill"),
        help="Action to run"
    )

    # parser.add_argument('--pid', type=int, default=42, help='what PID you want to kill')

    args = parser.parse_args()
    handle_receptor_action(args)