From bd7a0e44bdd8497dcae6582e40952eecd4022b7b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=A6=82=E5=A8=81?= Date: Sat, 23 Aug 2025 17:24:50 +0800 Subject: [PATCH] feat: cmd client --- backend/app/core/logger.py | 2 +- backend/app/main.py | 116 ++++++++++++++++++++++----- backend/app/nodes/cmd_node.py | 147 ++++++++++++++++++++++++++++------ backend/app/nodes/git_node.py | 17 ++++ 4 files changed, 237 insertions(+), 45 deletions(-) create mode 100644 backend/app/nodes/git_node.py diff --git a/backend/app/core/logger.py b/backend/app/core/logger.py index 93b1d20..3a7af28 100644 --- a/backend/app/core/logger.py +++ b/backend/app/core/logger.py @@ -27,7 +27,7 @@ ch = logging.StreamHandler() ch.setLevel(logging.DEBUG) # 设置带颜色的 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) logger.addHandler(ch) diff --git a/backend/app/main.py b/backend/app/main.py index ab01303..4f05d38 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -21,9 +21,11 @@ scheduler = AsyncIOScheduler() # 定时清理过期session def clear_session(): - TTL = 60 * 60 * 24 # 1天有效期 + TTL = 60 * 60 * 24 # 1天有效期 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: sessions.pop(token, None) @@ -31,7 +33,7 @@ def clear_session(): 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: raise HTTPException(status_code=401, detail="用户未登陆") _, user = sessions[token] @@ -50,7 +52,9 @@ async def lifespan(app: FastAPI): # 数据库 if await is_init(): 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("初始化数据库") # 定时器 scheduler.add_job(clear_session, "interval", seconds=5) @@ -92,7 +96,7 @@ def create_app(): logger.info("init_system") try: init_config = await load_init_config() - if init_config.get('is_init'): + if init_config.get("is_init"): return InitStatusResponse(data=True) # 保存配置 init_config["db_type"] = request.db_type @@ -104,7 +108,9 @@ def create_app(): init_config["admin_password"] = request.admin_password 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: await save_init_config(init_config) @@ -114,33 +120,103 @@ def create_app(): except Exception as 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): - 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: token = uuid.uuid4().hex sessions[token] = time.time(), user logger.info(sessions) - return LoginResponse(data={ - "token": token, - "user": user.json(del_columns=["password"]) - }) + return LoginResponse( + data={"token": token, "user": user.json(del_columns=["password"])} + ) return LoginResponse(success=False, message="帐号或密码错误") - @app.post("/logout", tags=["User"], description="退出登录", response_model=BaseResponse) - async def logout(token:str): + @app.post( + "/logout", tags=["User"], description="退出登录", response_model=BaseResponse + ) + async def logout(token: str): sessions.pop(token, None) logger.info(sessions) 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): - cmd_node = F.CMDNode() - cmd_node.set_params({"cmd": "ls"}) + nodes = [] - flow = F.AsyncFlow(cmd_node) - flow.set_params({"cmd": """ls && pwd"""}) - await flow.run_async({}) + local_node = F.CmdNode() + local_node.set_id("0") + 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() diff --git a/backend/app/nodes/cmd_node.py b/backend/app/nodes/cmd_node.py index 9877485..d428378 100644 --- a/backend/app/nodes/cmd_node.py +++ b/backend/app/nodes/cmd_node.py @@ -5,45 +5,144 @@ import asyncssh 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 不为空,则通过 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( 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: - # 远程执行 - conn = await asyncssh.connect( - host, - username=user, - password=password, - known_hosts=None, # 不检查 known_hosts - ) - result = await conn.run(cmd, check=False) - await conn.close() - return result.exit_status, result.stdout, result.stderr + 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): +class CmdNode(AsyncNode): + def set_id(self, id): + self.id = id + async def prep_async(self, shared): - pass + 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): - cmd = self.params.get("cmd") - _, out, err = await async_run_cmd(cmd) - if out: - logger.info(out) - if err: - logger.error(err) - return + 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): - pass + shared[self.id] = exec_res + + if exec_res.get("err"): + return "false" + return "true" diff --git a/backend/app/nodes/git_node.py b/backend/app/nodes/git_node.py new file mode 100644 index 0000000..c5a8bea --- /dev/null +++ b/backend/app/nodes/git_node.py @@ -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"