feat: cmd client
This commit is contained in:
parent
b8055a03e2
commit
bd7a0e44bd
|
|
@ -27,7 +27,7 @@ ch = logging.StreamHandler()
|
||||||
ch.setLevel(logging.DEBUG)
|
ch.setLevel(logging.DEBUG)
|
||||||
|
|
||||||
# 设置带颜色的 formatter
|
# 设置带颜色的 formatter
|
||||||
formatter = ColorFormatter("%(levelname)s: \t %(name)s - %(asctime)s - %(message)s")
|
formatter = ColorFormatter("%(levelname)s: \t %(name)s - %(asctime)s\n%(message)s")
|
||||||
ch.setFormatter(formatter)
|
ch.setFormatter(formatter)
|
||||||
|
|
||||||
logger.addHandler(ch)
|
logger.addHandler(ch)
|
||||||
|
|
|
||||||
|
|
@ -23,7 +23,9 @@ scheduler = AsyncIOScheduler()
|
||||||
def clear_session():
|
def clear_session():
|
||||||
TTL = 60 * 60 * 24 # 1天有效期
|
TTL = 60 * 60 * 24 # 1天有效期
|
||||||
now = time.time()
|
now = time.time()
|
||||||
expired_keys = [token for token, (expires_at, _) in sessions.items() if expires_at + TTL < now]
|
expired_keys = [
|
||||||
|
token for token, (expires_at, _) in sessions.items() if expires_at + TTL < now
|
||||||
|
]
|
||||||
for token in expired_keys:
|
for token in expired_keys:
|
||||||
sessions.pop(token, None)
|
sessions.pop(token, None)
|
||||||
|
|
||||||
|
|
@ -31,7 +33,7 @@ def clear_session():
|
||||||
logger.info(f"清理过期 now: {now} session: {expired_keys}")
|
logger.info(f"清理过期 now: {now} session: {expired_keys}")
|
||||||
|
|
||||||
|
|
||||||
def get_current_user(token:str = Depends(oauth2_scheme)) -> DB.User | None:
|
def get_current_user(token: str = Depends(oauth2_scheme)) -> DB.User | None:
|
||||||
if not token or token not in sessions:
|
if not token or token not in sessions:
|
||||||
raise HTTPException(status_code=401, detail="用户未登陆")
|
raise HTTPException(status_code=401, detail="用户未登陆")
|
||||||
_, user = sessions[token]
|
_, user = sessions[token]
|
||||||
|
|
@ -50,7 +52,9 @@ async def lifespan(app: FastAPI):
|
||||||
# 数据库
|
# 数据库
|
||||||
if await is_init():
|
if await is_init():
|
||||||
init_config = await load_init_config()
|
init_config = await load_init_config()
|
||||||
await init_tortoise(init_config["db_type"], init_config, init_config.get("admin_password"))
|
await init_tortoise(
|
||||||
|
init_config["db_type"], init_config, init_config.get("admin_password")
|
||||||
|
)
|
||||||
logger.info("初始化数据库")
|
logger.info("初始化数据库")
|
||||||
# 定时器
|
# 定时器
|
||||||
scheduler.add_job(clear_session, "interval", seconds=5)
|
scheduler.add_job(clear_session, "interval", seconds=5)
|
||||||
|
|
@ -92,7 +96,7 @@ def create_app():
|
||||||
logger.info("init_system")
|
logger.info("init_system")
|
||||||
try:
|
try:
|
||||||
init_config = await load_init_config()
|
init_config = await load_init_config()
|
||||||
if init_config.get('is_init'):
|
if init_config.get("is_init"):
|
||||||
return InitStatusResponse(data=True)
|
return InitStatusResponse(data=True)
|
||||||
# 保存配置
|
# 保存配置
|
||||||
init_config["db_type"] = request.db_type
|
init_config["db_type"] = request.db_type
|
||||||
|
|
@ -104,7 +108,9 @@ def create_app():
|
||||||
init_config["admin_password"] = request.admin_password
|
init_config["admin_password"] = request.admin_password
|
||||||
init_config["is_init"] = True
|
init_config["is_init"] = True
|
||||||
|
|
||||||
init_success = await init_tortoise(request.db_type, init_config, request.admin_password)
|
init_success = await init_tortoise(
|
||||||
|
request.db_type, init_config, request.admin_password
|
||||||
|
)
|
||||||
|
|
||||||
if init_success:
|
if init_success:
|
||||||
await save_init_config(init_config)
|
await save_init_config(init_config)
|
||||||
|
|
@ -114,33 +120,103 @@ def create_app():
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
return InitStatusResponse(success=False, message=str(e))
|
return InitStatusResponse(success=False, message=str(e))
|
||||||
|
|
||||||
@app.post("/login", tags=["User"], description="用户登录", response_model=LoginResponse)
|
@app.post(
|
||||||
|
"/login", tags=["User"], description="用户登录", response_model=LoginResponse
|
||||||
|
)
|
||||||
async def login(login_info: LoginRequest):
|
async def login(login_info: LoginRequest):
|
||||||
user = await DB.User.filter(username=login_info.username, password=login_info.get_md5_password()).first()
|
user = await DB.User.filter(
|
||||||
|
username=login_info.username, password=login_info.get_md5_password()
|
||||||
|
).first()
|
||||||
if user:
|
if user:
|
||||||
token = uuid.uuid4().hex
|
token = uuid.uuid4().hex
|
||||||
sessions[token] = time.time(), user
|
sessions[token] = time.time(), user
|
||||||
logger.info(sessions)
|
logger.info(sessions)
|
||||||
return LoginResponse(data={
|
return LoginResponse(
|
||||||
"token": token,
|
data={"token": token, "user": user.json(del_columns=["password"])}
|
||||||
"user": user.json(del_columns=["password"])
|
)
|
||||||
})
|
|
||||||
return LoginResponse(success=False, message="帐号或密码错误")
|
return LoginResponse(success=False, message="帐号或密码错误")
|
||||||
|
|
||||||
@app.post("/logout", tags=["User"], description="退出登录", response_model=BaseResponse)
|
@app.post(
|
||||||
async def logout(token:str):
|
"/logout", tags=["User"], description="退出登录", response_model=BaseResponse
|
||||||
|
)
|
||||||
|
async def logout(token: str):
|
||||||
sessions.pop(token, None)
|
sessions.pop(token, None)
|
||||||
logger.info(sessions)
|
logger.info(sessions)
|
||||||
return BaseResponse()
|
return BaseResponse()
|
||||||
|
|
||||||
@app.post("/run", tags=["Flow"], description="执行命令", response_model=BaseResponse)
|
@app.post(
|
||||||
|
"/run", tags=["Flow"], description="执行命令", response_model=BaseResponse
|
||||||
|
)
|
||||||
async def flow_run(id: int):
|
async def flow_run(id: int):
|
||||||
cmd_node = F.CMDNode()
|
nodes = []
|
||||||
cmd_node.set_params({"cmd": "ls"})
|
|
||||||
|
|
||||||
flow = F.AsyncFlow(cmd_node)
|
local_node = F.CmdNode()
|
||||||
flow.set_params({"cmd": """ls && pwd"""})
|
local_node.set_id("0")
|
||||||
await flow.run_async({})
|
nodes.append(local_node)
|
||||||
|
|
||||||
|
a800_node = F.CmdNode()
|
||||||
|
a800_node.set_id("1")
|
||||||
|
nodes.append(a800_node)
|
||||||
|
|
||||||
|
# 设置依赖关系
|
||||||
|
# node0 - "true" >> node1 - "true" >> node2
|
||||||
|
for i in range(len(nodes) -1):
|
||||||
|
# nodes[i] - "error" >> nodes[i + 1]
|
||||||
|
pass
|
||||||
|
|
||||||
|
flow = F.AsyncFlow(a800_node)
|
||||||
|
flow.set_params(
|
||||||
|
{
|
||||||
|
"0": {
|
||||||
|
"cmd": """
|
||||||
|
echo "1) 前端 - 拉取新代码"
|
||||||
|
cd /Users/liruwei/Documents/code/project/dabai/dbpai-web
|
||||||
|
git stash
|
||||||
|
git pull
|
||||||
|
|
||||||
|
echo "2) 后端 - 拉取新代码"
|
||||||
|
cd /Users/liruwei/Documents/code/project/dabai/dbpai-server
|
||||||
|
git stash
|
||||||
|
|
||||||
|
echo "3) 后端 - 更新 dbpai-chat-widget.min.js"
|
||||||
|
cp /Users/liruwei/Documents/code/project/dabai/dbpai-web/public/dbpai-chat-widget.min.js ./src/static
|
||||||
|
git commit -am "feat: dbpai-chat-widget.min.js"
|
||||||
|
git push
|
||||||
|
|
||||||
|
echo "4) 打包"
|
||||||
|
cd ..
|
||||||
|
sh build-dbpai-jenkisn.sh all
|
||||||
|
cd dbpai-jenkins
|
||||||
|
git commit -am "build"
|
||||||
|
git push
|
||||||
|
|
||||||
|
echo "5) 还原"
|
||||||
|
cd /Users/liruwei/Documents/code/project/dabai/dbpai-web
|
||||||
|
git stash pop
|
||||||
|
cd /Users/liruwei/Documents/code/project/dabai/dbpai-server
|
||||||
|
git stash pop
|
||||||
|
"""
|
||||||
|
},
|
||||||
|
"1": {
|
||||||
|
"cmd": """
|
||||||
|
cd /data/PAI/dbpai-beta/dbpai-jenkins
|
||||||
|
git stash
|
||||||
|
git pull origin v3.4.2
|
||||||
|
git stash pop
|
||||||
|
docker logs ai-guard-server -n 50
|
||||||
|
""",
|
||||||
|
"host": "58.214.239.10",
|
||||||
|
"user": "dgx07",
|
||||||
|
"password": "1f=bHt/.7kRo:KIXq",
|
||||||
|
"port": 6378,
|
||||||
|
"sudo": True,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
shared = {}
|
||||||
|
await flow.run_async(shared)
|
||||||
|
|
||||||
|
logger.info(shared)
|
||||||
|
|
||||||
return BaseResponse()
|
return BaseResponse()
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -5,45 +5,144 @@ import asyncssh
|
||||||
|
|
||||||
|
|
||||||
async def async_run_cmd(
|
async def async_run_cmd(
|
||||||
cmd: str, host: str = None, user: str = None, password: str = None
|
cmd: str,
|
||||||
|
host: str = None,
|
||||||
|
user: str = None,
|
||||||
|
password: str = None,
|
||||||
|
port: int = None,
|
||||||
|
sudo: bool = False,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
异步执行命令
|
异步执行命令
|
||||||
- 如果 host=None,则在本地执行
|
- 如果 host=None,则在本地执行
|
||||||
- 如果 host 不为空,则通过 SSH 执行远程命令
|
- 如果 host 不为空,则通过 SSH 执行远程命令
|
||||||
"""
|
"""
|
||||||
if host is None:
|
|
||||||
|
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(
|
proc = await asyncio.create_subprocess_shell(
|
||||||
cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
|
cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
|
||||||
)
|
)
|
||||||
stdout, stderr = await proc.communicate()
|
|
||||||
return proc.returncode, stdout.decode(), stderr.decode()
|
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:
|
else:
|
||||||
# 远程执行
|
try:
|
||||||
conn = await asyncssh.connect(
|
conn = await asyncssh.connect(
|
||||||
host,
|
host,
|
||||||
username=user,
|
username=user,
|
||||||
password=password,
|
password=password,
|
||||||
known_hosts=None, # 不检查 known_hosts
|
port=port,
|
||||||
|
known_hosts=None,
|
||||||
)
|
)
|
||||||
result = await conn.run(cmd, check=False)
|
except (asyncssh.PermissionDenied, asyncssh.DisconnectError) as e:
|
||||||
await conn.close()
|
# 认证失败 / 连接断开
|
||||||
return result.exit_status, result.stdout, result.stderr
|
yield "stderr", f"SSH connection failed: {str(e)}"
|
||||||
|
return
|
||||||
|
except Exception as e:
|
||||||
class CMDNode(AsyncNode):
|
# 其它异常
|
||||||
async def prep_async(self, shared):
|
yield "stderr", f"SSH error: {str(e)}"
|
||||||
pass
|
|
||||||
|
|
||||||
async def exec_async(self, prep_res):
|
|
||||||
cmd = self.params.get("cmd")
|
|
||||||
_, out, err = await async_run_cmd(cmd)
|
|
||||||
if out:
|
|
||||||
logger.info(out)
|
|
||||||
if err:
|
|
||||||
logger.error(err)
|
|
||||||
return
|
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))
|
||||||
|
else:
|
||||||
|
out.append(str(line))
|
||||||
|
logger.info(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):
|
async def post_async(self, shared, prep_res, exec_res):
|
||||||
pass
|
shared[self.id] = exec_res
|
||||||
|
|
||||||
|
if exec_res.get("err"):
|
||||||
|
return "false"
|
||||||
|
return "true"
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,17 @@
|
||||||
|
from core.pocket_flow import AsyncNode
|
||||||
|
|
||||||
|
class GitNode(AsyncNode):
|
||||||
|
def set_id(self, id):
|
||||||
|
self.id = id
|
||||||
|
|
||||||
|
async def prep_async(self, shared):
|
||||||
|
pass
|
||||||
|
|
||||||
|
async def exec_async(self, prep_res):
|
||||||
|
return None
|
||||||
|
|
||||||
|
async def post_async(self, shared, prep_res, exec_res):
|
||||||
|
if exec_res:
|
||||||
|
return "true"
|
||||||
|
else:
|
||||||
|
return "false"
|
||||||
Loading…
Reference in New Issue