Reapply "Feat/run console commands"

This reverts commit 4cda54ab60.
This commit is contained in:
Josiah Baldwin
2025-09-27 14:48:55 -07:00
parent e0694f980c
commit 12a3040f89
8 changed files with 135 additions and 29 deletions

View File

@@ -184,7 +184,7 @@ class Session(object):
async def _listen_data_task(self, websocket):
async for message in websocket:
await self._eventer.emit("raw", message)
# Meshcentral does pong wrong and breaks our parsing, so fix it here.
# Meshcentral does pong wrong and breaks our parsing, so fix it here. This is fixed now, but we want compatibility with old versions.
if message == '{action:"pong"}':
message = '{"action":"pong"}'
@@ -1497,7 +1497,7 @@ class Session(object):
return nid
result = {n: {"complete": False, "result": [], "command": command} for n in nodeids}
async def _():
async def _console():
async for event in self.events({"action": "msg", "type": "console"}):
node = match_nodeid(event["nodeid"], nodeids)
if node:
@@ -1509,7 +1509,25 @@ class Session(object):
elif (event["value"].startswith("Run commands")):
continue
result[node]["result"].append(event["value"])
async def __(command):
# We create this task AFTER getting the first message, but I don't feel like implementing this twice, so we'll pass in the first message and have it parsed immediately
async def _reply(responseid, start_data=None):
# Returns True when all results are in, Falsey otherwise
def _parse_event(event):
node = match_nodeid(event["nodeid"], nodeids)
if node:
result.setdefault(node, {})["complete"] = True
result[node]["result"].append(event["result"])
if all(_["complete"] for key, _ in result.items()):
return True
if start_data is not None:
if _parse_event(start_data):
return
async for event in self.events({"action": "msg", "type": "runcommands", "responseid": responseid}):
if _parse_event(event):
break
async def __(command, tg, tasks):
data = await self._send_command(command, "run_command", timeout=timeout)
if data.get("type", None) != "runcommands" and data.get("result", "ok").lower() != "ok":
@@ -1544,9 +1562,81 @@ class Session(object):
tasks = []
async with asyncio.TaskGroup() as tg:
if expect_response:
tasks.append(tg.create_task(asyncio.wait_for(_(), timeout=timeout)))
tasks.append(tg.create_task(__({ "action": 'runcommands', "nodeids": nodeids, "type": (2 if powershell else 0), "cmds": command, "runAsUser": runAsUser })))
tasks.append(tg.create_task(__({ "action": 'runcommands', "nodeids": nodeids, "type": (2 if powershell else 0), "cmds": command, "runAsUser": runAsUser, "reply": not ignore_output}, tg, tasks)))
return {n: v | {"result": "".join(v["result"])} for n,v in result.items()}
async def run_console_command(self, nodeids, command, powershell=False, runasuser=False, runasuseronly=False, ignore_output=False, timeout=None):
'''
Run a mesh console command on any number of nodes. WARNING: Non namespaced call. Calling this function again before it returns may cause unintended consequences.
Args:
nodeids (str|list[str]): Unique ids of nodes on which to run the command
command (str): Command to run
ignore_output (bool): Don't bother trying to get the output. Every device will return an empty string for its result.
timeout (int): duration in seconds to wait for a response before throwing an error
Returns:
dict[str, ~meshctrl.types.RunCommandResponse]: Dict containing mapped output of the commands by device
Raises:
:py:class:`~meshctrl.exceptions.ServerError`: Error text from server if there is a failure
:py:class:`~meshctrl.exceptions.SocketError`: Info about socket closure
ValueError: `Invalid device id` if device is not found
asyncio.TimeoutError: Command timed out
'''
if isinstance(nodeids, str):
nodeids = [nodeids]
def match_nodeid(id, ids):
for nid in ids:
if (nid == id):
return nid
if (nid[6:] == id):
return nid
if (f"node//{nid}" == id):
return nid
result = {n: {"complete": False, "result": [], "command": command} for n in nodeids}
async def _console():
async for event in self.events({"action": "msg", "type": "console"}):
node = match_nodeid(event["nodeid"], nodeids)
if node:
result[node]["result"].append(event["value"])
result.setdefault(node, {})["complete"] = True
if all(_["complete"] for key, _ in result.items()):
break
async def __(command, tg, tasks):
data = await self._send_command(command, "run_console_command", timeout=timeout)
if data.get("type", None) != "runcommands" and data.get("result", "ok").lower() != "ok":
raise exceptions.ServerError(data["result"])
elif data.get("type", None) != "runcommands" and data.get("result", "ok").lower() == "ok":
expect_response = False
console_task = tg.create_task(asyncio.wait_for(_console(), timeout=timeout))
if not ignore_output:
userid = (await self.user_info())["_id"]
for n in nodeids:
device_info = await self.device_info(n, timeout=timeout)
try:
permissions = device_info.mesh.links.get(userid, {}).get("rights",constants.DeviceRights.norights)\
# This should work for device rights, but it only seems to work for mesh rights. Not sure why, but I can't get the events to show up when the user only has individual device rights
# |device_info.get("links", {}).get(userid, {}).get("rights", constants.DeviceRights.norights)
# If we don't have agentconsole rights, we won't be able te read the output, so fill in blanks on this node
if not permissions&constants.DeviceRights.agentconsole:
result[n]["complete"] = True
else:
expect_response = True
except AttributeError:
result[n]["complete"] = True
if expect_response:
tasks.append(console_task)
else:
console_task.cancel()
tasks = []
async with asyncio.TaskGroup() as tg:
tasks.append(tg.create_task(__({ "action": 'runcommands', "nodeids": nodeids, "type": 4, "cmds": command}, tg, tasks)))
return {n: v | {"result": "".join(v["result"])} for n,v in result.items()}