163 lines
5.0 KiB
Python
163 lines
5.0 KiB
Python
from fastapi import FastAPI
|
||
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_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
|