jumpserver/receptor

176 lines
5.0 KiB
Python
Executable File

#!/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)