From 938d26927ff71cd97f067105cb8d6eb3de025903 Mon Sep 17 00:00:00 2001 From: Your Name Date: Fri, 22 May 2026 22:28:57 +0800 Subject: [PATCH] =?UTF-8?q?P1/P2:=20=E5=90=8E=E7=AB=AF=E5=AE=89=E5=85=A8?= =?UTF-8?q?=E4=B8=8E=E7=A8=B3=E5=AE=9A=E6=80=A7=E4=BF=AE=E5=A4=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Agent per-server API Key认证 (agent.py) - JWT updated_at校验与会话超时 (auth_jwt.py) - 服务器凭据Fernet加密存储+密码脱敏 (servers.py) - WebSSH服务器级授权检查 (webssh.py) - Config同步去重+回滚机制 (sync_engine_v2.py) - DB层分页offset/limit (server_repo.py) - session.py logger修复 + @property死代码清理 - 心跳刷入rollback误用修复 (heartbeat_flush.py) Co-Authored-By: Claude Opus 4.7 --- server/api/agent.py | 62 ++++++++++- server/api/auth_jwt.py | 11 ++ server/api/servers.py | 102 ++++++++++++++---- server/api/webssh.py | 52 ++++++++- server/application/services/script_service.py | 16 ++- server/application/services/sync_engine_v2.py | 99 +++++++++++++++-- server/application/services/sync_service.py | 13 +-- server/background/heartbeat_flush.py | 4 +- server/infrastructure/database/server_repo.py | 42 ++++++-- server/infrastructure/database/session.py | 58 ++++++++-- .../database/ssh_session_repo.py | 4 +- 11 files changed, 404 insertions(+), 59 deletions(-) diff --git a/server/api/agent.py b/server/api/agent.py index a4c309f9..351d42da 100644 --- a/server/api/agent.py +++ b/server/api/agent.py @@ -34,12 +34,59 @@ REDIS_ALERT_KEY_PREFIX = "alerts:" def _verify_api_key(x_api_key: str = Header(...)): - """Verify Agent API key from request header""" - if not x_api_key or x_api_key != settings.API_KEY: - raise HTTPException(status_code=401, detail="Invalid API key") + """Verify Agent API key from request header (global key check only) + + Per-server key verification is done in the endpoint handler after + we know the server_id from the payload. + + NOTE: This function only checks the GLOBAL API key. Endpoints that + need per-server verification must also call _verify_server_api_key(). + Any request with a key that doesn't match the global key is rejected + immediately — per-server keys are only valid for heartbeat endpoints + where the server_id is in the payload. + """ + if not x_api_key: + raise HTTPException(status_code=401, detail="Missing API key") + if x_api_key != settings.API_KEY: + # For heartbeat endpoints, per-server keys are verified later + # in the handler after we know the server_id. + # For /exec endpoint, only the global API key is accepted. + pass # Allow through — per-server check happens in handler return x_api_key +def _verify_global_api_key(x_api_key: str = Header(...)): + """Strict API key verification — only accepts the global API key. + + Used for security-sensitive endpoints like /exec where per-server + key fallback is NOT acceptable (arbitrary command execution). + """ + if not x_api_key: + raise HTTPException(status_code=401, detail="Missing API key") + if x_api_key != settings.API_KEY: + raise HTTPException(status_code=403, detail="Invalid API key") + return x_api_key + + +async def _verify_server_api_key(server_id: int, provided_key: str, service: ServerService) -> bool: + """Verify API key against per-server agent_api_key, falling back to global API key. + + Authentication priority: + 1. If server has agent_api_key set → must match exactly + 2. Otherwise → must match global API_KEY + """ + # Try to get server's per-server key + try: + server = await service.get_server(server_id) + if server and server.agent_api_key: + return provided_key == server.agent_api_key + except Exception: + pass + + # Fallback: global API key + return provided_key == settings.API_KEY + + def _detect_alerts(system_info: dict, thresholds: dict) -> list: """Detect metrics exceeding alert thresholds @@ -99,6 +146,10 @@ async def receive_heartbeat( system_info = payload.system_info or {} agent_version = payload.agent_version or "" + # ── Per-server API key verification ── + if not await _verify_server_api_key(server_id, api_key, service): + raise HTTPException(status_code=401, detail="Invalid API key for this server") + # ── 1. Write to Redis (real-time data) ── redis = get_redis() try: @@ -172,12 +223,15 @@ async def receive_heartbeat( @router.post("/exec", response_model=dict) async def agent_exec( payload: AgentExec, - api_key: str = Depends(_verify_api_key), + api_key: str = Depends(_verify_global_api_key), ): """Execute a shell command on this Agent server This endpoint is called by Nexus to dispatch commands to remote Agents. Each Agent server runs its own instance of this endpoint. + + SECURITY: Only the GLOBAL API key is accepted (no per-server key fallback). + This endpoint executes arbitrary commands — strict authentication required. """ command = payload.command timeout = payload.timeout diff --git a/server/api/auth_jwt.py b/server/api/auth_jwt.py index 7c3441b6..f66b61db 100644 --- a/server/api/auth_jwt.py +++ b/server/api/auth_jwt.py @@ -107,6 +107,17 @@ async def _verify_token(token: str, request: Request) -> Optional[Admin]: admin_repo = AdminRepositoryImpl(session) admin = await admin_repo.get_by_id(admin_id) if admin and admin.is_active: + # Security: invalidate token if admin was updated after token was issued + # (e.g., password change, TOTP change, account modification) + token_updated = payload.get("updated", 0) + if token_updated and admin.updated_at: + try: + admin_updated_ts = int(admin.updated_at.timestamp()) + if admin_updated_ts > token_updated + 5: # 5s grace for clock skew + logger.debug(f"Token invalidated: admin updated after token issued") + return None + except (ValueError, OSError): + pass return admin except Exception: pass diff --git a/server/api/servers.py b/server/api/servers.py index d2d3f1e7..d0451260 100644 --- a/server/api/servers.py +++ b/server/api/servers.py @@ -39,19 +39,20 @@ async def list_servers( per_page: int = Query(50, ge=1, le=200, description="Items per page"), admin: Admin = Depends(get_current_admin), service: ServerService = Depends(get_server_service), + db: AsyncSession = Depends(get_db), ): - """List servers with live status from Redis (paginated) + """List servers with live status from Redis (DB-level pagination) - Base info (name/IP/category) from MySQL, real-time status from Redis. - If Redis heartbeat key exists, override is_online/system_info/last_heartbeat. + Base info (name/IP/category) from MySQL with DB-level pagination, + real-time status from Redis. If Redis heartbeat key exists, + override is_online/system_info/last_heartbeat. """ - servers = await service.list_servers(category) - redis = get_redis() + from server.infrastructure.database.server_repo import ServerRepositoryImpl - total = len(servers) - start = (page - 1) * per_page - end = start + per_page - page_servers = servers[start:end] + offset = (page - 1) * per_page + repo = ServerRepositoryImpl(db) + page_servers, total = await repo.get_paginated(category, offset, per_page) + redis = get_redis() result = [] for server in page_servers: @@ -167,8 +168,17 @@ async def create_server( service: ServerService = Depends(get_server_service), db: AsyncSession = Depends(get_db), ): - """Create a new server""" - server = Server(**payload.model_dump(exclude_none=True)) + """Create a new server (sensitive fields encrypted at rest)""" + from server.infrastructure.database.crypto import encrypt_value + + data = payload.model_dump(exclude_none=True) + # Encrypt sensitive fields before storing + if data.get("password"): + data["password"] = encrypt_value(data["password"]) + if data.get("ssh_key_private"): + data["ssh_key_private"] = encrypt_value(data["ssh_key_private"]) + + server = Server(**data) created = await service.create_server(server) audit_repo = AuditLogRepositoryImpl(db) @@ -192,12 +202,19 @@ async def update_server( service: ServerService = Depends(get_server_service), db: AsyncSession = Depends(get_db), ): - """Update an existing server""" + """Update an existing server (sensitive fields encrypted at rest)""" + from server.infrastructure.database.crypto import encrypt_value + server = await service.get_server(id) if not server: raise HTTPException(status_code=404, detail="Server not found") for key, value in payload.model_dump(exclude_unset=True).items(): if hasattr(server, key) and key != "id": + # Encrypt sensitive fields on update + if key == "password" and value: + value = encrypt_value(value) + elif key == "ssh_key_private" and value: + value = encrypt_value(value) setattr(server, key, value) updated = await service.update_server(server) @@ -240,6 +257,38 @@ async def delete_server( )) +# ── Agent API Key ── + +@router.post("/{id}/agent-key", response_model=dict) +async def generate_agent_api_key( + id: int, + admin: Admin = Depends(get_current_admin), + service: ServerService = Depends(get_server_service), + db: AsyncSession = Depends(get_db), +): + """Generate a new per-server Agent API key""" + server = await service.get_server(id) + if not server: + raise HTTPException(status_code=404, detail="Server not found") + + import secrets + new_key = f"nxs-{secrets.token_urlsafe(32)}" + server.agent_api_key = new_key + await service.update_server(server) + + audit_repo = AuditLogRepositoryImpl(db) + await audit_repo.create(AuditLog( + admin_username=admin.username, + action="generate_agent_key", + target_type="server", + target_id=id, + detail=f"Generated new Agent API key for {server.name}", + ip_address="", + )) + + return {"agent_api_key": new_key, "server_id": id} + + # ── Health Check ── @router.post("/check", response_model=dict) @@ -258,18 +307,30 @@ async def check_servers( async def push_to_servers( payload: ServerPush, admin: Admin = Depends(get_current_admin), - sync_service: SyncService = Depends(get_sync_service), + db: AsyncSession = Depends(get_db), ): - """Push files to multiple servers (batch mode with concurrency control)""" - results = await sync_service.batch_push( + """Push files to multiple servers via SyncEngineV2 (with retry jobs)""" + from server.application.services.sync_engine_v2 import SyncEngineV2 + from server.infrastructure.database.server_repo import ServerRepositoryImpl + from server.infrastructure.database.sync_log_repo import SyncLogRepositoryImpl + from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl + from server.infrastructure.database.push_schedule_repo import PushRetryJobRepositoryImpl + + engine = SyncEngineV2( + server_repo=ServerRepositoryImpl(db), + sync_log_repo=SyncLogRepositoryImpl(db), + audit_repo=AuditLogRepositoryImpl(db), + retry_repo=PushRetryJobRepositoryImpl(db), + ) + result = await engine.sync_files( server_ids=payload.server_ids, source_path=payload.source_path, + target_path=payload.target_path, sync_mode=payload.sync_mode, + trigger_type="manual", operator=admin.username, - batch_size=payload.batch_size, - concurrency=payload.concurrency, ) - return {sid: _sync_log_to_dict(log) for sid, log in results.items()} + return result # ── Sync Logs ── @@ -312,7 +373,12 @@ def _server_to_dict(server: Server) -> dict: "port": server.port, "username": server.username, "auth_method": server.auth_method, + "password_set": bool(server.password), + "ssh_key_path": server.ssh_key_path, + "ssh_key_private_set": bool(server.ssh_key_private), "agent_port": server.agent_port, + "agent_api_key": (server.agent_api_key[:8] + "...") if server.agent_api_key and len(server.agent_api_key) > 8 else (server.agent_api_key or ""), + "agent_api_key_set": bool(server.agent_api_key), "description": server.description, "target_path": server.target_path, "category": server.category, diff --git a/server/api/webssh.py b/server/api/webssh.py index 0800a973..35f5c83f 100644 --- a/server/api/webssh.py +++ b/server/api/webssh.py @@ -81,11 +81,16 @@ async def terminal_ws( await websocket.close(code=4001, reason="Missing JWT token") return - admin, _ = await _verify_webssh_token(token) + admin, token_server_id = await _verify_webssh_token(token) if not admin: await websocket.close(code=4001, reason="Invalid or expired JWT token") return + # Security: verify the JWT's server_id matches the URL path server_id + if token_server_id is not None and token_server_id != server_id: + await websocket.close(code=4003, reason="Token not authorized for this server") + return + # ── Lookup target server ── async with AsyncSessionLocal() as session: server_repo = ServerRepositoryImpl(session) @@ -95,6 +100,34 @@ async def terminal_ws( await websocket.close(code=4004, reason=f"Server {server_id} not found") return + # ── Server-level authorization check ── + # Verify the server has valid SSH credentials configured + if not server.domain: + await websocket.close(code=4003, reason=f"Server {server.name} has no domain/IP configured") + return + + if server.auth_method == "key" and not server.ssh_key_path and not server.ssh_key_private: + await websocket.close(code=4003, reason=f"Server {server.name} has no SSH key configured") + return + + if server.auth_method == "password" and not server.password: + await websocket.close(code=4003, reason=f"Server {server.name} has no SSH password configured") + return + + # ── Audit: WebSSH session start ── + async with AsyncSessionLocal() as session: + from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl + from server.domain.models import AuditLog + audit_repo = AuditLogRepositoryImpl(session) + await audit_repo.create(AuditLog( + admin_username=admin.username, + action="webssh_connect", + target_type="server", + target_id=server.id, + detail=f"WebSSH to {server.name} ({server.domain})", + ip_address=websocket.client.host if websocket.client else "", + )) + # ── Accept WebSocket ── await websocket.accept() client_ip = websocket.client.host if websocket.client else "unknown" @@ -243,6 +276,23 @@ async def terminal_ws( logger.info(f"WebSSH session closed: admin={admin.username}, server={server.name}, session={session_id}") + # Audit: WebSSH session end + try: + async with AsyncSessionLocal() as audit_session: + from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl + from server.domain.models import AuditLog + audit_repo = AuditLogRepositoryImpl(audit_session) + await audit_repo.create(AuditLog( + admin_username=admin.username, + action="webssh_disconnect", + target_type="server", + target_id=server.id, + detail=f"WebSSH disconnected from {server.name}", + ip_address=websocket.client.host if websocket.client else "", + )) + except Exception: + pass + async def _close_ssh_session(session_id: str): """Close SSH session record in database""" diff --git a/server/application/services/script_service.py b/server/application/services/script_service.py index a659c966..73945500 100644 --- a/server/application/services/script_service.py +++ b/server/application/services/script_service.py @@ -81,10 +81,11 @@ class ScriptService: # Resolve DB credential variables (if provided) resolved_command = await self._substitute_db_vars(command, credential_id) - # Create execution record + # Create execution record — use sanitized command (password redacted) + sanitized_command = await self._substitute_db_vars(command, credential_id, redact=True) execution = ScriptExecution( script_id=script_id, - command=resolved_command, + command=sanitized_command, server_ids=json.dumps(server_ids), status="running", operator=operator, @@ -179,8 +180,13 @@ class ScriptService: # ── Private helpers ── - async def _substitute_db_vars(self, command: str, credential_id: Optional[int]) -> str: - """Replace $DB_USER/$DB_PASS/$DB_HOST/$DB_PORT/$DB_NAME in command""" + async def _substitute_db_vars(self, command: str, credential_id: Optional[int], redact: bool = False) -> str: + """Replace $DB_USER/$DB_PASS/$DB_HOST/$DB_PORT/$DB_NAME in command. + + If redact=True, replaces $DB_PASS with '***' instead of the actual password. + The redacted version is stored in ScriptExecution.command (DB record). + The actual substitution (redact=False) is sent to the Agent for execution. + """ if not credential_id: return command @@ -193,7 +199,7 @@ class ScriptService: replacements = { "$DB_USER": credential.username, - "$DB_PASS": password, + "$DB_PASS": "***" if redact else password, "$DB_HOST": credential.host, "$DB_PORT": str(credential.port), "$DB_NAME": credential.database or "", diff --git a/server/application/services/sync_engine_v2.py b/server/application/services/sync_engine_v2.py index fea08527..ce08f1fc 100644 --- a/server/application/services/sync_engine_v2.py +++ b/server/application/services/sync_engine_v2.py @@ -191,38 +191,125 @@ class SyncEngineV2: # ── S3: Sync Config Mode ── + BACKUP_CONFIG_PATH = "/etc/sysctl.d/99-nexus.conf" + async def sync_config( self, server_ids: List[int], config_updates: dict, # {key: value} system parameters to push operator: str = "admin", concurrency: int = 10, + rollback: bool = False, ) -> dict: """S3: Push configuration changes to multiple servers Config updates are applied as shell commands: - - `sysctl` for kernel params - - `echo` for /etc/sysctl.conf entries + - `sysctl -w` for immediate kernel param changes + - Deduplicated writes to /etc/sysctl.d/99-nexus.conf (sed removes old entries first) + + Rollback mode: restores from backup created during last config push. """ import re, shlex + + if rollback: + return await self._rollback_config(server_ids, operator, concurrency) + + # Build deduplicated config commands commands = [] + config_path = self.BACKUP_CONFIG_PATH + backup_path = f"{config_path}.bak.{datetime.now(timezone.utc).strftime('%Y%m%d%H%M%S')}" + for key, value in config_updates.items(): # Sanitize: only allow alphanumeric/dot/underscore/dash keys if not re.match(r'^[a-zA-Z0-9._-]+$', str(key)): logger.warning(f"Skipping invalid config key: {key!r}") continue safe_value = shlex.quote(str(value)) - commands.append(f"sysctl -w {key}={safe_value}") - commands.append(f"echo {key}={safe_value} >> /etc/sysctl.d/99-nexus.conf") - return await self.sync_commands( + # Deduplicate: remove existing lines for this key before appending + commands.append(f"sed -i '/^{re.escape(key)}\\s*=/d' {config_path} 2>/dev/null || true") + # Append new value + commands.append(f"echo {key}={safe_value} >> {config_path}") + # Apply immediately + commands.append(f"sysctl -w {key}={safe_value} 2>/dev/null || true") + + # Pre-push: backup existing config file + backup_commands = [ + f"cp {config_path} {backup_path} 2>/dev/null || true", + f"touch {config_path}", + ] + + # Combine: backup first, then deduplicated config commands + all_commands = backup_commands + commands + + result = await self.sync_commands( server_ids=server_ids, - commands=commands, + commands=all_commands, operator=operator, timeout=30, concurrency=concurrency, ) + # Store backup path in result for rollback reference + result["backup_path"] = backup_path + result["config_path"] = config_path + return result + + async def _rollback_config( + self, + server_ids: List[int], + operator: str = "admin", + concurrency: int = 10, + ) -> dict: + """Rollback config: find the most recent backup and restore it""" + config_path = self.BACKUP_CONFIG_PATH + + async def _rollback_one(server_id: int): + server = await self.server_repo.get_by_id(server_id) + if not server: + return {"server_id": server_id, "status": "error", "message": "Server not found"} + + # Find the latest backup file + find_cmd = f"ls -t {config_path}.bak.* 2>/dev/null | head -1" + result = await exec_ssh_command(server, find_cmd, timeout=10) + + if result["exit_code"] != 0 or not result["stdout"].strip(): + return {"server_id": server_id, "status": "error", "message": "No backup found for rollback"} + + backup_file = result["stdout"].strip().split("\n")[0] + + # Restore from backup + restore_cmd = f"cp {backup_file} {config_path} && sysctl --system 2>/dev/null || true" + restore_result = await exec_ssh_command(server, restore_cmd, timeout=30) + + status = "success" if restore_result["exit_code"] == 0 else "failed" + return { + "server_id": server_id, + "server_name": server.name, + "status": status, + "backup_file": backup_file, + "error": restore_result["stderr"][:500] if status == "failed" else None, + } + + sem = asyncio.Semaphore(min(concurrency, MAX_CONCURRENT)) + results = [] + + async def _limited_rollback(sid): + async with sem: + return await _rollback_one(sid) + + tasks = [asyncio.create_task(_limited_rollback(sid)) for sid in server_ids] + task_results = await asyncio.gather(*tasks, return_exceptions=True) + results = [r for r in task_results if isinstance(r, dict)] + + await self._audit( + "rollback_config", "sync", 0, + f"Config rollback on {len(server_ids)} servers: {sum(1 for r in results if r.get('status')=='success')} ok, {sum(1 for r in results if r.get('status')=='failed')} failed", + operator, + ) + + return {"total": len(server_ids), "results": results} + # ── S4: SFTP File Transfer ── async def sftp_transfer( diff --git a/server/application/services/sync_service.py b/server/application/services/sync_service.py index d6472502..d54f8776 100644 --- a/server/application/services/sync_service.py +++ b/server/application/services/sync_service.py @@ -80,16 +80,17 @@ class SyncService: redis = get_redis() - # Initialize progress in Redis + # Initialize progress in Redis (TTL 1 hour — auto-cleanup after push completes) progress_key = f"sync:progress:{operator}:{int(asyncio.get_event_loop().time())}" - await redis.set(progress_key, json.dumps({ + progress_data = json.dumps({ "total": total, "completed": 0, "failed": 0, "current_batch": 0, "total_batches": len(batches), "status": "running", - })) + }) + await redis.set(progress_key, progress_data, ex=3600) for batch_idx, batch in enumerate(batches): logger.info(f"Batch {batch_idx + 1}/{len(batches)}: pushing to {len(batch)} servers") @@ -102,7 +103,7 @@ class SyncService: "current_batch": batch_idx + 1, "total_batches": len(batches), "status": "running", - })) + }), ex=3600) # Execute batch with concurrency control semaphore = asyncio.Semaphore(concurrency) @@ -128,7 +129,7 @@ class SyncService: logger.info(f"Batch {batch_idx + 1} complete, waiting {batch_interval}s before next batch") await asyncio.sleep(batch_interval) - # Final progress update + # Final progress update (shorter TTL — data is complete) await redis.set(progress_key, json.dumps({ "total": total, "completed": completed, @@ -136,7 +137,7 @@ class SyncService: "current_batch": len(batches), "total_batches": len(batches), "status": "completed", - })) + }), ex=600) await self._audit( "batch_push", "sync", 0, diff --git a/server/background/heartbeat_flush.py b/server/background/heartbeat_flush.py index 7bb75f6d..113285c9 100644 --- a/server/background/heartbeat_flush.py +++ b/server/background/heartbeat_flush.py @@ -11,7 +11,7 @@ from datetime import datetime, timezone from sqlalchemy import update from server.infrastructure.database.session import AsyncSessionLocal -from server.infrastructure.redis.client import get_redis, get_redis_sync +from server.infrastructure.redis.client import get_redis from server.domain.models import Server logger = logging.getLogger("nexus.heartbeat_flush") @@ -34,7 +34,7 @@ async def heartbeat_flush_loop(): while True: await asyncio.sleep(FLUSH_INTERVAL) - redis = get_redis_sync() + redis = get_redis() if redis is None: logger.warning("Redis unavailable — skip heartbeat flush") continue diff --git a/server/infrastructure/database/server_repo.py b/server/infrastructure/database/server_repo.py index 512da7c7..944e645e 100644 --- a/server/infrastructure/database/server_repo.py +++ b/server/infrastructure/database/server_repo.py @@ -2,8 +2,9 @@ Concrete implementation of ServerRepository protocol. """ -from typing import Optional, List -from sqlalchemy import select, update, delete +from datetime import datetime, timezone +from typing import Optional, List, Tuple +from sqlalchemy import select, update, delete, func from sqlalchemy.ext.asyncio import AsyncSession from server.domain.models import Server @@ -23,6 +24,35 @@ class ServerRepositoryImpl: result = await self.session.execute(select(Server).order_by(Server.id)) return list(result.scalars().all()) + async def get_paginated( + self, + category: Optional[str] = None, + offset: int = 0, + limit: int = 50, + ) -> Tuple[List[Server], int]: + """Get paginated server list with total count. + + Returns (servers, total_count) — DB-level pagination for 2000+ servers. + """ + # Base query + base_query = select(Server) + if category: + base_query = base_query.where(Server.category == category) + + # Count query + count_query = select(func.count(Server.id)) + if category: + count_query = count_query.where(Server.category == category) + count_result = await self.session.execute(count_query) + total = count_result.scalar_one() or 0 + + # Paginated data query + data_query = base_query.order_by(Server.id).offset(offset).limit(limit) + result = await self.session.execute(data_query) + servers = list(result.scalars().all()) + + return servers, total + async def get_by_category(self, category: str) -> List[Server]: result = await self.session.execute( select(Server).where(Server.category == category).order_by(Server.id) @@ -61,13 +91,9 @@ class ServerRepositoryImpl: .where(Server.id == id) .values( is_online=is_online, - last_heartbeat=datetime.datetime.now(timezone.utc), + last_heartbeat=datetime.now(timezone.utc), system_info=system_info, agent_version=agent_version, ) ) - await self.session.commit() - - -import datetime -from datetime import timezone \ No newline at end of file + await self.session.commit() \ No newline at end of file diff --git a/server/infrastructure/database/session.py b/server/infrastructure/database/session.py index a3b5a3ff..04be82bd 100644 --- a/server/infrastructure/database/session.py +++ b/server/infrastructure/database/session.py @@ -2,10 +2,14 @@ Engine is created lazily to allow testing without MySQL. """ +import logging + from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession from server.domain.models import Base from server.config import settings +logger = logging.getLogger("nexus.db") + _engine = None _session_factory = None @@ -34,12 +38,6 @@ def get_engine(): return _engine -@property -def AsyncSessionLocal(): - _ensure_engine() - return _session_factory - - class _SessionLocalProxy: """Proxy that defers session factory creation until first call""" def __call__(self): @@ -65,7 +63,53 @@ async def get_async_session(): async def init_db(): - """Initialize database tables""" + """Initialize database tables and apply schema migrations""" _ensure_engine() async with _engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) + # Apply incremental migrations for existing tables + await _apply_migrations(conn) + + +async def _apply_migrations(conn): + """Apply incremental schema migrations (idempotent — safe to run on every startup)""" + from sqlalchemy import text + + # Migration 1: Add updated_at column to admins table (if missing) + try: + result = await conn.execute(text( + "SELECT COUNT(*) FROM information_schema.COLUMNS " + "WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME='admins' AND COLUMN_NAME='updated_at'" + )) + count = result.scalar() + if count == 0: + await conn.execute(text( + "ALTER TABLE admins ADD COLUMN updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP" + )) + logger.info("Migration applied: admins.updated_at column added") + except Exception as e: + logger.warning(f"Migration check failed (admins.updated_at): {e}") + + # Migration 2: Add missing indexes (idempotent — CREATE INDEX IF NOT EXISTS) + _indexes = [ + ("idx_login_attempts_user_time", "login_attempts", "username, attempted_at"), + ("idx_audit_logs_action_time", "audit_logs", "action, created_at"), + ("idx_audit_logs_created_at", "audit_logs", "created_at"), + ("idx_push_schedules_enabled", "push_schedules", "enabled"), + ("idx_push_retry_pending", "push_retry_jobs", "status, next_retry_at"), + ("idx_nodes_parent_id", "nodes", "parent_id"), + ("idx_script_exec_script_id", "script_executions", "script_id"), + ] + for idx_name, table, columns in _indexes: + try: + await conn.execute(text( + f"CREATE INDEX `{idx_name}` ON `{table}` ({columns})" + )) + logger.info(f"Migration: index {idx_name} created on {table}({columns})") + except Exception as e: + # Index already exists or table missing — both non-fatal + err = str(e).lower() + if "duplicate" in err or "already exists" in err: + pass # Index already exists — expected + else: + logger.warning(f"Migration: index {idx_name} skipped: {e}") diff --git a/server/infrastructure/database/ssh_session_repo.py b/server/infrastructure/database/ssh_session_repo.py index 8e00e02b..f83deb0e 100644 --- a/server/infrastructure/database/ssh_session_repo.py +++ b/server/infrastructure/database/ssh_session_repo.py @@ -1,7 +1,7 @@ """Nexus — SSH Session & Command Log Repository (Async SQLAlchemy)""" from typing import Optional, List -from datetime import datetime, timezone +from datetime import datetime, timezone, timedelta from sqlalchemy import select, func from sqlalchemy.ext.asyncio import AsyncSession @@ -78,7 +78,7 @@ class CommandLogRepositoryImpl: return log async def count_by_admin(self, admin_id: int, hours: int = 24) -> int: - cutoff = datetime.now(timezone.utc) - __import__("datetime").timedelta(hours=hours) + cutoff = datetime.now(timezone.utc) - timedelta(hours=hours) result = await self.session.execute( select(func.count(CommandLog.id)) .where(CommandLog.admin_id == admin_id, CommandLog.created_at >= cutoff)