from core.pocket_flow import AsyncNode from core.logger import logger import asyncio import asyncssh async def async_run_cmd( cmd: str, host: str = None, user: str = None, password: str = None, port: int = None, sudo: bool = False, ): """ 异步执行命令 - 如果 host=None,则在本地执行 - 如果 host 不为空,则通过 SSH 执行远程命令 """ async def read_stream(stream, stream_type): async for line in stream: text = line if not isinstance(text, str): text = line.decode() yield stream_type, text.rstrip() if not host: # 本地执行 proc = await asyncio.create_subprocess_shell( cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) tasks = [ read_stream(proc.stdout, "stdout"), read_stream(proc.stderr, "stderr"), ] for t in tasks: async for stream_type, line in t: yield stream_type, line await proc.wait() else: try: conn = await asyncssh.connect( host, username=user, password=password, port=port, known_hosts=None, ) except (asyncssh.PermissionDenied, asyncssh.DisconnectError) as e: # 认证失败 / 连接断开 yield "stderr", f"SSH connection failed: {str(e)}" return except Exception as e: # 其它异常 yield "stderr", f"SSH error: {str(e)}" return try: # 执行脚本 remote_cmd = cmd if sudo: remote_cmd = f"echo {password!r} | sudo -S -p '' bash -c '{cmd}'" # 执行 proc = await conn.create_process(remote_cmd, term_type="xterm") # 分别创建 stdout/stderr 的异步生成器 stdout_gen = read_stream(proc.stdout, "stdout") stderr_gen = read_stream(proc.stderr, "stderr") # 用 task 管理两个流 stdout_task = asyncio.create_task(stdout_gen.__anext__()) stderr_task = asyncio.create_task(stderr_gen.__anext__()) while True: if stdout_task is None and stderr_task is None: break # 等待任意一个流先返回 done, _ = await asyncio.wait( [t for t in (stdout_task, stderr_task) if t is not None], return_when=asyncio.FIRST_COMPLETED, ) for fut in done: try: stream_type, line = fut.result() yield stream_type, line # 继续拉取该流的下一行 if stream_type == "stdout": stdout_task = asyncio.create_task(stdout_gen.__anext__()) else: stderr_task = asyncio.create_task(stderr_gen.__anext__()) except StopAsyncIteration: # 该流已结束 if fut == stdout_task: stdout_task = None elif fut == stderr_task: stderr_task = None await proc.wait() finally: conn.close() await conn.wait_closed() class CmdNode(AsyncNode): def set_id(self, id): self.id = id async def prep_async(self, shared): 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): 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 if exec_res.get("err"): return "false" return "true"