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