your-workflow/backend/app/nodes/cmd_node.py

59 lines
1.7 KiB
Python

from fastapi import FastAPI
from core.pocket_flow import AsyncNode
from core.logger import logger
from .base import async_run_cmd
class CmdNode(AsyncNode):
def set_info(self, id, callback):
self.id = id
self.callback = callback
async def prep_async(self, shared):
if self.callback:
await self.callback(self.id, "exec")
return {
"cmd": self.params[self.id].get("cmd"),
"host": self.params[self.id].get("host"),
"user": self.params[self.id].get("user"),
"password": self.params[self.id].get("password"),
"port": self.params[self.id].get("port"),
"sudo": self.params[self.id].get("sudo", False),
}
async def exec_async(self, prep_res):
if self.callback:
await self.callback(self.id, "exec")
err = []
out = []
async for stream_type, line in async_run_cmd(**prep_res):
if stream_type == "stderr":
err.append(str(line))
logger.debug(f'err: {line}')
else:
out.append(str(line))
logger.debug(f"out: {line}")
err_str = "\n".join(err)
out_str = "\n".join(out)
# 补充一下密码失效问题
if "Sorry, try again." in out_str and not err_str:
err_str = "SSH connection failed no password was provided"
return {"err": err_str, "out": out_str}
async def post_async(self, shared, prep_res, exec_res):
shared[self.id] = exec_res
result = "true"
if exec_res.get("err"):
result = "false"
if self.callback:
await self.callback(self.id, "post")
return result