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, wait: int = 0, ): """ 异步执行命令 - 如果 host=None,则在本地执行 - 如果 host 不为空,则通过 SSH 执行远程命令 """ if wait > 0: await asyncio.sleep(wait) 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()