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

150 lines
4.7 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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"