106 lines
3.4 KiB
Python
106 lines
3.4 KiB
Python
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()
|