move run script to nats
This commit is contained in:
parent
4780af910c
commit
e4b1f39fdc
|
@ -436,7 +436,7 @@ class Agent(BaseAuditModel):
|
|||
except Exception:
|
||||
return "err"
|
||||
|
||||
async def nats_cmd(self, data, timeout=30):
|
||||
async def nats_cmd(self, data, timeout=30, wait=True):
|
||||
nc = NATS()
|
||||
options = {
|
||||
"servers": f"tls://{settings.ALLOWED_HOSTS[0]}:4222",
|
||||
|
@ -450,15 +450,22 @@ class Agent(BaseAuditModel):
|
|||
except:
|
||||
return "natsdown"
|
||||
|
||||
try:
|
||||
msg = await nc.request(self.agent_id, msgpack.dumps(data), timeout=timeout)
|
||||
except ErrTimeout:
|
||||
ret = "timeout"
|
||||
else:
|
||||
ret = msgpack.loads(msg.data)
|
||||
if wait:
|
||||
try:
|
||||
msg = await nc.request(
|
||||
self.agent_id, msgpack.dumps(data), timeout=timeout
|
||||
)
|
||||
except ErrTimeout:
|
||||
ret = "timeout"
|
||||
else:
|
||||
ret = msgpack.loads(msg.data)
|
||||
|
||||
await nc.close()
|
||||
return ret
|
||||
await nc.close()
|
||||
return ret
|
||||
else:
|
||||
await nc.publish(self.agent_id, msgpack.dumps(data))
|
||||
await nc.flush()
|
||||
await nc.close()
|
||||
|
||||
def salt_api_cmd(self, **kwargs):
|
||||
|
||||
|
|
|
@ -38,7 +38,7 @@ from winupdate.serializers import WinUpdatePolicySerializer
|
|||
|
||||
from .tasks import uninstall_agent_task, send_agent_update_task
|
||||
from winupdate.tasks import bulk_check_for_updates_task
|
||||
from scripts.tasks import run_script_bg_task, run_bulk_script_task
|
||||
from scripts.tasks import run_bulk_script_task
|
||||
|
||||
from tacticalrmm.utils import notify_error
|
||||
|
||||
|
@ -660,10 +660,7 @@ def recover(request):
|
|||
def run_script(request):
|
||||
agent = get_object_or_404(Agent, pk=request.data["pk"])
|
||||
script = get_object_or_404(Script, pk=request.data["scriptPK"])
|
||||
|
||||
output = request.data["output"]
|
||||
args = request.data["args"]
|
||||
|
||||
req_timeout = int(request.data["timeout"]) + 3
|
||||
|
||||
AuditLog.audit_script_run(
|
||||
|
@ -672,48 +669,21 @@ def run_script(request):
|
|||
script=script.name,
|
||||
)
|
||||
|
||||
data = {
|
||||
"func": "runscript",
|
||||
"timeout": request.data["timeout"],
|
||||
"script_args": request.data["args"],
|
||||
"payload": {
|
||||
"code": script.code,
|
||||
"shell": script.shell,
|
||||
},
|
||||
}
|
||||
|
||||
if output == "wait":
|
||||
r = agent.salt_api_cmd(
|
||||
timeout=req_timeout,
|
||||
func="win_agent.run_script",
|
||||
kwargs={
|
||||
"filepath": script.filepath,
|
||||
"filename": script.filename,
|
||||
"shell": script.shell,
|
||||
"timeout": request.data["timeout"],
|
||||
"args": args,
|
||||
},
|
||||
)
|
||||
|
||||
if isinstance(r, dict):
|
||||
if r["stdout"]:
|
||||
return Response(r["stdout"])
|
||||
elif r["stderr"]:
|
||||
return Response(r["stderr"])
|
||||
else:
|
||||
try:
|
||||
r["retcode"]
|
||||
except KeyError:
|
||||
return notify_error("Something went wrong")
|
||||
|
||||
return Response(f"Return code: {r['retcode']}")
|
||||
|
||||
else:
|
||||
if r == "timeout":
|
||||
return notify_error("Unable to contact the agent")
|
||||
elif r == "error":
|
||||
return notify_error("Something went wrong")
|
||||
else:
|
||||
return notify_error(str(r))
|
||||
|
||||
r = asyncio.run(agent.nats_cmd(data, timeout=req_timeout))
|
||||
return Response(r)
|
||||
else:
|
||||
data = {
|
||||
"agentpk": agent.pk,
|
||||
"scriptpk": script.pk,
|
||||
"timeout": request.data["timeout"],
|
||||
"args": args,
|
||||
}
|
||||
run_script_bg_task.delay(data)
|
||||
asyncio.run(agent.nats_cmd(data, wait=False))
|
||||
return Response(f"{script.name} will now be run on {agent.hostname}")
|
||||
|
||||
|
||||
|
|
|
@ -3,23 +3,6 @@ from agents.models import Agent
|
|||
from .models import Script
|
||||
|
||||
|
||||
@app.task
|
||||
def run_script_bg_task(data):
|
||||
agent = Agent.objects.get(pk=data["agentpk"])
|
||||
script = Script.objects.get(pk=data["scriptpk"])
|
||||
|
||||
agent.salt_api_async(
|
||||
func="win_agent.run_script",
|
||||
kwargs={
|
||||
"filepath": script.filepath,
|
||||
"filename": script.filename,
|
||||
"shell": script.shell,
|
||||
"timeout": data["timeout"],
|
||||
"args": data["args"],
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@app.task
|
||||
def run_bulk_script_task(data):
|
||||
# for powershell and batch scripts only, workaround for salt bg script bug
|
||||
|
|
Loading…
Reference in New Issue