mirror of https://github.com/jumpserver/jumpserver
176 lines
5.0 KiB
Python
Executable File
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)
|