mirror of https://github.com/jumpserver/jumpserver
145 lines
4.8 KiB
Python
145 lines
4.8 KiB
Python
import concurrent.futures
|
|
import os
|
|
import queue
|
|
import socket
|
|
|
|
from django.conf import settings
|
|
import ansible_runner
|
|
from receptorctl import ReceptorControl
|
|
|
|
from ops.ansible.cleaner import WorkPostRunCleaner, cleanup_post_run
|
|
|
|
|
|
class ReceptorCtl:
|
|
@property
|
|
def ctl(self):
|
|
return ReceptorControl(settings.ANSIBLE_RECEPTOR_SOCK_PATH)
|
|
|
|
def cancel(self, unit_id):
|
|
return self.ctl.simple_command("work cancel {}".format(unit_id))
|
|
|
|
def nodes(self):
|
|
return self.ctl.simple_command("status").get("Advertisements", None)
|
|
|
|
def submit_work(self,
|
|
worktype,
|
|
payload,
|
|
node=None,
|
|
tlsclient=None,
|
|
ttl=None,
|
|
signwork=False,
|
|
params=None, ):
|
|
return self.ctl.submit_work(worktype, payload, node, tlsclient, ttl, signwork, params)
|
|
|
|
def get_work_results(self, unit_id, startpos=0, return_socket=False, return_sockfile=True):
|
|
return self.ctl.get_work_results(unit_id, startpos, return_socket, return_sockfile)
|
|
|
|
def kill_process(self, pid):
|
|
submit_result = self.submit_work(worktype="kill", node="primary", payload=str(pid))
|
|
unit_id = submit_result["unitid"]
|
|
result_socket, result_file = self.get_work_results(unit_id=unit_id, return_sockfile=True,
|
|
return_socket=True)
|
|
while not result_socket.close():
|
|
buf = result_file.read()
|
|
if not buf:
|
|
break
|
|
print(buf.decode('utf8'))
|
|
|
|
|
|
receptor_ctl = ReceptorCtl()
|
|
|
|
|
|
def run(**kwargs):
|
|
receptor_runner = AnsibleReceptorRunner(**kwargs)
|
|
return receptor_runner.run()
|
|
|
|
|
|
class AnsibleReceptorRunner(WorkPostRunCleaner):
|
|
def __init__(self, **kwargs):
|
|
self.runner_params = kwargs
|
|
self.unit_id = None
|
|
|
|
def write_unit_id(self):
|
|
if not self.unit_id:
|
|
return
|
|
private_dir = self.runner_params.get("private_data_dir", "")
|
|
with open(os.path.join(private_dir, "local.unitid"), "w") as f:
|
|
f.write(self.unit_id)
|
|
f.flush()
|
|
|
|
@property
|
|
def clean_dir(self):
|
|
return self.runner_params.get("private_data_dir", None)
|
|
|
|
@cleanup_post_run
|
|
def run(self):
|
|
input, output = socket.socketpair()
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
|
|
transmitter_future = executor.submit(self.transmit, input)
|
|
result = receptor_ctl.submit_work(payload=output.makefile('rb'),
|
|
node='primary', worktype='ansible-runner')
|
|
input.close()
|
|
output.close()
|
|
|
|
self.unit_id = result['unitid']
|
|
self.write_unit_id()
|
|
|
|
transmitter_future.result()
|
|
|
|
result_file = receptor_ctl.get_work_results(self.unit_id, return_sockfile=True)
|
|
|
|
stdout_queue = queue.Queue()
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
|
|
processor_future = executor.submit(self.processor, result_file, stdout_queue)
|
|
|
|
while not processor_future.done() or \
|
|
not stdout_queue.empty():
|
|
msg = stdout_queue.get()
|
|
if msg is None:
|
|
break
|
|
print(msg)
|
|
|
|
return processor_future.result()
|
|
|
|
def transmit(self, _socket):
|
|
try:
|
|
ansible_runner.run(
|
|
streamer='transmit',
|
|
_output=_socket.makefile('wb'),
|
|
**self.runner_params
|
|
)
|
|
finally:
|
|
_socket.shutdown(socket.SHUT_WR)
|
|
|
|
def processor(self, _result_file, stdout_queue):
|
|
try:
|
|
original_event_handler = self.runner_params.pop("event_handler", None)
|
|
original_status_handler = self.runner_params.pop("status_handler", None)
|
|
|
|
def event_handler(data, **kwargs):
|
|
stdout = data.get('stdout', '')
|
|
if stdout:
|
|
stdout_queue.put(stdout)
|
|
if original_event_handler:
|
|
original_event_handler(data, **kwargs)
|
|
|
|
def status_handler(data, **kwargs):
|
|
private_data_dir = self.runner_params.get("private_data_dir", None)
|
|
if private_data_dir:
|
|
data["private_data_dir"] = private_data_dir
|
|
if original_status_handler:
|
|
original_status_handler(data, **kwargs)
|
|
|
|
return ansible_runner.interface.run(
|
|
quite=True,
|
|
streamer='process',
|
|
_input=_result_file,
|
|
event_handler=event_handler,
|
|
status_handler=status_handler,
|
|
**self.runner_params,
|
|
)
|
|
finally:
|
|
stdout_queue.put(None)
|