devops-copilot 0.1.0__py3-none-any.whl

app/modules/chat/router.py

@@ -0,0 +1,362 @@
+import jwt
+from fastapi import APIRouter, Depends, HTTPException, status, WebSocket, WebSocketDisconnect, Query
+from typing import List
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlmodel import select
+from langchain_core.messages import HumanMessage, AIMessage
+
+from app.core.config import settings
+from app.core.database import get_session, async_session_maker
+from app.modules.auth import User
+from app.modules.auth.service import get_current_user
+from app.modules.servers.models import Server
+from app.modules.servers.service import servers_service
+from app.modules.chat.models import ChatSession, ChatMessage, AgentAction
+from app.modules.chat.schema import (
+    ChatSessionCreate,
+    ChatSessionResponse,
+    ChatMessageResponse,
+    AgentActionResponse
+)
+from app.modules.chat.agent import (
+    create_agent_executor,
+    StreamingCallbackHandler,
+    active_websocket,
+    active_session_id,
+    active_owner_id,
+    decrypt_data,
+    asyncssh
+)
+
+router = APIRouter()
+
+# --- REST Endpoints for Chat History ---
+
+@router.post("/sessions", response_model=ChatSessionResponse, status_code=status.HTTP_201_CREATED)
+async def create_session(
+    payload: ChatSessionCreate,
+    current_user: User = Depends(get_current_user),
+    session: AsyncSession = Depends(get_session)
+):
+    """Create a new chat session thread for the logged-in administrator."""
+    new_session = ChatSession(title=payload.title, user_id=current_user.id)
+    session.add(new_session)
+    await session.commit()
+    await session.refresh(new_session)
+    return new_session
+
+@router.get("/sessions", response_model=List[ChatSessionResponse])
+async def list_sessions(
+    current_user: User = Depends(get_current_user),
+    session: AsyncSession = Depends(get_session)
+):
+    """Retrieve all chat sessions belonging to the current administrator."""
+    statement = select(ChatSession).where(ChatSession.user_id == current_user.id).order_by(ChatSession.created_at.desc())
+    result = await session.exec(statement)
+    return result.all()
+
+@router.get("/sessions/{session_id}/messages", response_model=List[ChatMessageResponse])
+async def get_messages(
+    session_id: int,
+    current_user: User = Depends(get_current_user),
+    session: AsyncSession = Depends(get_session)
+):
+    """Retrieve historical messages for a specific chat session thread."""
+    # Verify session ownership
+    sess_statement = select(ChatSession).where(ChatSession.id == session_id, ChatSession.user_id == current_user.id)
+    sess_result = await session.exec(sess_statement)
+    if not sess_result.first():
+        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail="Chat session not found")
+
+    msg_statement = select(ChatMessage).where(ChatMessage.session_id == session_id).order_by(ChatMessage.created_at.asc())
+    msg_result = await session.exec(msg_statement)
+    return msg_result.all()
+
+@router.delete("/sessions/{session_id}")
+async def delete_session(
+    session_id: int,
+    current_user: User = Depends(get_current_user),
+    session: AsyncSession = Depends(get_session)
+):
+    """Delete a chat session and all its associated messages and actions."""
+    # Verify session ownership
+    sess_statement = select(ChatSession).where(ChatSession.id == session_id, ChatSession.user_id == current_user.id)
+    sess_result = await session.exec(sess_statement)
+    chat_session = sess_result.first()
+    if not chat_session:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="Chat session not found"
+        )
+
+    # Delete messages
+    msg_statement = select(ChatMessage).where(ChatMessage.session_id == session_id)
+    msg_result = await session.exec(msg_statement)
+    for msg in msg_result.all():
+        await session.delete(msg)
+
+    # Delete actions
+    act_statement = select(AgentAction).where(AgentAction.session_id == session_id)
+    act_result = await session.exec(act_statement)
+    for act in act_result.all():
+        await session.delete(act)
+
+    await session.delete(chat_session)
+    await session.commit()
+    return {"status": "success", "message": f"Chat session {session_id} deleted successfully"}
+
+# --- WebSocket Real-Time Chat & Execution Tunnel ---
+
+@router.websocket("/ws")
+async def websocket_endpoint(
+    websocket: WebSocket,
+    token: str = Query(...)
+):
+    """
+    Establish a WebSocket connection for real-time DevOps chat and live execution logging.
+    Validates the admin's JWT token query parameter prior to connection acceptance.
+    """
+    # 1. Manual JWT Token Verification for Handshake
+    try:
+        payload = jwt.decode(token, settings.JWT_SECRET_KEY, algorithms=[settings.JWT_ALGORITHM])
+        username = payload.get("sub")
+        if not username:
+            await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+            return
+
+        async with async_session_maker() as session:
+            statement = select(User).where(User.username == username)
+            result = await session.exec(statement)
+            user = result.first()
+            if not user or not user.is_active:
+                await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+                return
+    except Exception:
+        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+        return
+
+    # Accept client WebSocket connection
+    await websocket.accept()
+
+    try:
+        while True:
+            # Receive inbound JSON payloads
+            data = await websocket.receive_json()
+            msg_type = data.get("type")
+
+            if msg_type == "message":
+                session_id = data.get("session_id")
+                content = data.get("content")
+
+                if not session_id or not content:
+                    continue
+
+                # 2. Load conversation history for memory injection
+                async with async_session_maker() as db_session:
+                    # Verify session ownership
+                    sess_stmt = select(ChatSession).where(ChatSession.id == session_id, ChatSession.user_id == user.id)
+                    sess = await db_session.exec(sess_stmt)
+                    if not sess.first():
+                        await websocket.send_json({"type": "error", "data": "Unauthorized session access"})
+                        continue
+
+                    # Record user message in DB
+                    user_msg = ChatMessage(session_id=session_id, sender="user", content=content)
+                    db_session.add(user_msg)
+                    await db_session.commit()
+
+                    # Fetch history
+                    history_stmt = select(ChatMessage).where(ChatMessage.session_id == session_id).order_by(ChatMessage.created_at.asc())
+                    history_result = await db_session.exec(history_stmt)
+                    db_messages = history_result.all()
+
+                chat_history = []
+                # Map history database records to LangChain Human/AI message formats
+                for msg in db_messages[:-1]:  # Exclude current message since it is input
+                    if msg.sender == "user":
+                        chat_history.append(HumanMessage(content=msg.content))
+                    else:
+                        chat_history.append(AIMessage(content=msg.content))
+
+                # 3. Setup context variables for the SSH execution tools
+                ws_token = active_websocket.set(websocket)
+                sess_token = active_session_id.set(session_id)
+                owner_token = active_owner_id.set(user.id)
+
+                try:
+                    # Stream tokens via custom callback handler
+                    callback = StreamingCallbackHandler(websocket)
+                    agent_executor = create_agent_executor(callbacks=[callback])
+
+                    # Invoke agent compiled StateGraph using messages schema
+                    inputs = {
+                        "messages": chat_history + [HumanMessage(content=content)]
+                    }
+                    agent_response = await agent_executor.ainvoke(inputs)
+
+                    # Extract final text from output message list
+                    ai_text = agent_response["messages"][-1].content
+                    
+                    # Record final agent response in database
+                    async with async_session_maker() as db_session:
+                        ai_msg = ChatMessage(session_id=session_id, sender="ai", content=ai_text)
+                        db_session.add(ai_msg)
+                        await db_session.commit()
+
+                    # Notify client execution complete
+                    await websocket.send_json({
+                        "type": "finished",
+                        "data": ai_text
+                    })
+
+                except Exception as e:
+                    await websocket.send_json({"type": "error", "data": f"Agent Execution Error: {str(e)}"})
+                finally:
+                    # Clean context variables
+                    active_websocket.reset(ws_token)
+                    active_session_id.reset(sess_token)
+                    active_owner_id.reset(owner_token)
+
+            elif msg_type == "approve" or msg_type == "reject":
+                action_id = data.get("action_id")
+                if not action_id:
+                    continue
+
+                async with async_session_maker() as db_session:
+                    # Fetch queued action
+                    act_stmt = select(AgentAction).where(AgentAction.id == action_id)
+                    act_result = await db_session.exec(act_stmt)
+                    action = act_result.first()
+
+                    if not action or action.status != "pending":
+                        await websocket.send_json({"type": "error", "data": "Action not found or already processed"})
+                        continue
+
+                    # Verify session ownership
+                    sess_stmt = select(ChatSession).where(ChatSession.id == action.session_id, ChatSession.user_id == user.id)
+                    sess = await db_session.exec(sess_stmt)
+                    if not sess.first():
+                        await websocket.send_json({"type": "error", "data": "Unauthorized session access"})
+                        continue
+
+                    if msg_type == "reject":
+                        action.status = "rejected"
+                        await db_session.commit()
+                        await websocket.send_json({
+                            "type": "stdout",
+                            "data": f"\n[Action {action_id} Rejected by User]\n"
+                        })
+                        continue
+
+                    # If approved, update status
+                    action.status = "approved"
+                    await db_session.commit()
+
+                    # Load targeted server config
+                    server_stmt = select(Server).where(Server.id == action.server_id)
+                    srv_result = await db_session.exec(server_stmt)
+                    server = srv_result.first()
+
+                await websocket.send_json({
+                    "type": "stdout",
+                    "data": f"\n[Action {action_id} Approved] Running command on {server.name}...\n"
+                })
+
+                # 4. Decrypt and execute approved write command
+                credential = decrypt_data(server.encrypted_credential)
+                conn_args = {
+                    "host": server.ip_address,
+                    "port": server.ssh_port,
+                    "username": server.ssh_username,
+                    "known_hosts": None
+                }
+                if server.ssh_auth_type == "password":
+                    conn_args["password"] = credential
+                else:
+                    conn_args["client_keys"] = [asyncssh.import_private_key(credential)]
+
+                stdout_accumulated = []
+                stderr_accumulated = []
+                exit_status = -1
+
+                try:
+                    # Connect and execute process, streaming stdout/stderr line-by-line
+                    async with asyncssh.connect(**conn_args) as conn:
+                        async with conn.create_process(action.command) as process:
+                            async for line in process.stdout:
+                                if isinstance(line, bytes):
+                                    line = line.decode("utf-8", errors="replace")
+                                stdout_accumulated.append(line)
+                                await websocket.send_json({"type": "stdout", "data": line})
+                            async for line in process.stderr:
+                                if isinstance(line, bytes):
+                                    line = line.decode("utf-8", errors="replace")
+                                stderr_accumulated.append(line)
+                                await websocket.send_json({"type": "stderr", "data": line})
+                            exit_status = await process.wait()
+                            exit_code = exit_status.exit_status if exit_status.exit_status is not None else 0
+
+                    await websocket.send_json({
+                        "type": "stdout",
+                        "data": f"[Process finished with code {exit_code}]\n"
+                    })
+
+                    stdout_str = "".join(stdout_accumulated)
+                    stderr_str = "".join(stderr_accumulated)
+
+                    # Auto-index approved command output into RAG knowledge base
+                    try:
+                        from app.modules.knowledge.service import index_command_output
+                        index_command_output(server.name, action.command, stdout_str, stderr_str, exit_code)
+                    except Exception:
+                        pass
+
+                    # 5. Invoke agent loop with the execution output
+                    ws_token = active_websocket.set(websocket)
+                    sess_token = active_session_id.set(action.session_id)
+                    owner_token = active_owner_id.set(user.id)
+
+                    try:
+                        callback = StreamingCallbackHandler(websocket)
+                        agent_executor = create_agent_executor(callbacks=[callback])
+
+                        # Feed output back to agent reasoning chain
+                        resume_prompt = (
+                            f"User approved execution of action '{action_id}'.\n"
+                            f"Command '{action.command}' executed on server '{server.name}'.\n"
+                            f"Exit Code: {exit_code}\n"
+                            f"Stdout:\n{stdout_str}\n"
+                            f"Stderr:\n{stderr_str}\n\n"
+                            f"Please analyze these outputs and provide your final response to the user."
+                        )
+
+                        # Invoke the agent graph with the resume message
+                        agent_response = await agent_executor.ainvoke({
+                            "messages": [HumanMessage(content=resume_prompt)]
+                        })
+
+                        ai_text = agent_response["messages"][-1].content
+                        
+                        # Save final summary response in database
+                        async with async_session_maker() as db_session:
+                            ai_msg = ChatMessage(session_id=action.session_id, sender="ai", content=ai_text)
+                            db_session.add(ai_msg)
+                            await db_session.commit()
+
+                        await websocket.send_json({
+                            "type": "finished",
+                            "data": ai_text
+                        })
+
+                    except Exception as e:
+                        await websocket.send_json({"type": "error", "data": f"Error during agent completion: {str(e)}"})
+                    finally:
+                        active_websocket.reset(ws_token)
+                        active_session_id.reset(sess_token)
+                        active_owner_id.reset(owner_token)
+
+                except Exception as e:
+                    await websocket.send_json({"type": "error", "data": f"SSH Execution Failed: {str(e)}"})
+
+    except WebSocketDisconnect:
+        pass

app/modules/chat/schema.py

@@ -0,0 +1,39 @@
+from datetime import datetime
+from pydantic import BaseModel, Field
+
+class ChatSessionCreate(BaseModel):
+    """Payload to instantiate a new chat session."""
+    title: str = Field(..., min_length=1, max_length=100)
+
+class ChatSessionResponse(BaseModel):
+    """Metadata response representing a chat session."""
+    id: int
+    title: str
+    user_id: int
+    created_at: datetime
+
+    class Config:
+        from_attributes = True
+
+class ChatMessageResponse(BaseModel):
+    """Metadata response representing a single message in a session thread."""
+    id: int
+    session_id: int
+    sender: str
+    content: str
+    created_at: datetime
+
+    class Config:
+        from_attributes = True
+
+class AgentActionResponse(BaseModel):
+    """Metadata response representing a queued server command approval."""
+    id: str
+    session_id: int
+    server_id: int
+    command: str
+    status: str
+    created_at: datetime
+
+    class Config:
+        from_attributes = True

app/modules/guardrails/__init__.py

@@ -0,0 +1 @@
+# This makes the rag folder a Python package

app/modules/guardrails/schema.py

@@ -0,0 +1,22 @@
+from pydantic import BaseModel, Field
+
+class SecurityCheckRequest(BaseModel):
+    """Input payload to check a proposed bash command for security threats."""
+    command: str = Field(..., min_length=1)
+
+class SecurityCheckResponse(BaseModel):
+    """Result of the semantic safety check against the dangerous commands database."""
+    safe: bool
+    reason: str
+    similarity: float
+
+class BlacklistPatternCreate(BaseModel):
+    """Payload to register a new forbidden pattern inside the semantic guardrail database."""
+    pattern: str = Field(..., min_length=3)
+    reason: str = Field(..., min_length=3)
+
+class BlacklistPatternResponse(BaseModel):
+    """Details of a registered dangerous command pattern in the vector store."""
+    id: str
+    pattern: str
+    reason: str

app/modules/guardrails/service.py

@@ -0,0 +1,121 @@
+import uuid
+import chromadb
+from typing import List, Tuple
+
+# Initialize ChromaDB persistent client pointing to data volume
+# Will resolve to /app/data/chroma_db inside Docker containers
+client = chromadb.PersistentClient(path="data/chroma_db")
+
+# Fetch or create the security guardrail collection
+collection = client.get_or_create_collection(name="security_blacklist")
+
+async def init_and_seed_db():
+    """
+    Checks if the security blacklist collection has documents.
+    If empty, populates it with default dangerous command templates and reasons.
+    """
+    if collection.count() == 0:
+        default_seeds = [
+            {
+                "pattern": "rm -rf /",
+                "reason": "Recursive deletion of the root directory. Results in absolute OS loss."
+            },
+            {
+                "pattern": "rm -rf /etc",
+                "reason": "Recursive deletion of configuration directories. Breaks system configuration."
+            },
+            {
+                "pattern": "mkfs.ext4 /dev/sda",
+                "reason": "Formatting raw storage drive. Destroys filesystem structure and partition data."
+            },
+            {
+                "pattern": "dd if=/dev/zero of=/dev/sda",
+                "reason": "Zero-writing raw block device. Wipes partition tables and master boot records."
+            },
+            {
+                "pattern": "iptables -F",
+                "reason": "Flushing active firewall rules. Exposes the server to external connection attempts."
+            },
+            {
+                "pattern": "cat /etc/shadow",
+                "reason": "Exposing system shadow user file. Leaks hashed credentials database."
+            },
+            {
+                "pattern": "chmod 777 -R /",
+                "reason": "Recursive permissive access modifier. Allows unprivileged write and execution globally."
+            },
+            {
+                "pattern": "reboot",
+                "reason": "Forces hard system restart. Terminates running containers and processes instantly."
+            },
+            {
+                "pattern": "shutdown -h now",
+                "reason": "Forces host system shutdown. Takes the server offline."
+            }
+        ]
+
+        ids = [f"seed_{i}" for i in range(len(default_seeds))]
+        documents = [item["pattern"] for item in default_seeds]
+        metadatas = [{"reason": item["reason"]} for item in default_seeds]
+
+        # Sync add to collection
+        collection.add(
+            ids=ids,
+            documents=documents,
+            metadatas=metadatas
+        )
+
+def is_command_safe(command: str) -> Tuple[bool, str, float]:
+    """
+    Performs semantic similarity check of proposed shell command against blacklist.
+    Returns (is_safe, reason, similarity_score).
+    """
+    # Fetch closest match (n_results=1)
+    results = collection.query(
+        query_texts=[command],
+        n_results=1
+    )
+
+    if results and results["documents"] and len(results["documents"][0]) > 0:
+        matched_doc = results["documents"][0][0]
+        distance = results["distances"][0][0]  # L2 distance metric
+        metadata = results["metadatas"][0][0]
+
+        # Convert L2 distance to Cosine Similarity score: 1.0 - (L2 / 2.0)
+        # ChromaDB ONNX embeddings are L2-normalized, making this conversion mathematically valid.
+        similarity = 1.0 - (distance / 2.0)
+
+        # Flag commands with a similarity index higher than 0.75 (L2 distance < 0.5)
+        if similarity > 0.75:
+            return (
+                False,
+                f"Blocked: Command matches dangerous pattern '{matched_doc}'. Reason: {metadata['reason']}",
+                similarity
+            )
+
+        return True, "No critical security threats detected.", similarity
+
+    return True, "Guardrail threat database is empty.", 0.0
+
+def add_blacklist_pattern(pattern: str, reason: str) -> dict:
+    """Adds a new custom forbidden pattern to the ChromaDB vector database."""
+    pattern_id = f"custom_{uuid.uuid4().hex}"
+    collection.add(
+        ids=[pattern_id],
+        documents=[pattern],
+        metadatas=[{"reason": reason}]
+    )
+    return {"id": pattern_id, "pattern": pattern, "reason": reason}
+
+def get_all_blacklist_patterns() -> List[dict]:
+    """Retrieves all registered dangerous patterns inside the RAG collection."""
+    results = collection.get()
+    patterns = []
+    if results and "ids" in results:
+        for i in range(len(results["ids"])):
+            patterns.append({
+                "id": results["ids"][i],
+                "pattern": results["documents"][i],
+                "reason": results["metadatas"][i]["reason"]
+            })
+    return patterns

app/modules/knowledge/__init__.py

@@ -0,0 +1 @@
+# This makes the knowledge folder a Python package

app/modules/knowledge/service.py

@@ -0,0 +1,123 @@
+import uuid
+import chromadb
+from typing import List, Optional
+
+# Shared ChromaDB persistent client (same data volume as guardrails)
+client = chromadb.PersistentClient(path="data/chroma_db")
+
+# --- 3 Separate Collections ---
+command_history = client.get_or_create_collection(name="command_history")
+server_logs = client.get_or_create_collection(name="server_logs")
+server_configs = client.get_or_create_collection(name="server_configs")
+
+CHUNK_SIZE = 500
+CHUNK_OVERLAP = 50
+
+def _chunk_text(text: str) -> List[str]:
+    """Split text into overlapping chunks for better retrieval."""
+    if len(text) <= CHUNK_SIZE:
+        return [text]
+    chunks = []
+    start = 0
+    while start < len(text):
+        end = start + CHUNK_SIZE
+        chunks.append(text[start:end])
+        start = end - CHUNK_OVERLAP
+    return chunks
+
+
+def index_command_output(
+    server_name: str, command: str, stdout: str, stderr: str, exit_code: int
+):
+    """Auto-index SSH command execution results into command_history collection."""
+    doc = (
+        f"Server: {server_name}\n"
+        f"Command: {command}\n"
+        f"Exit Code: {exit_code}\n"
+        f"Stdout:\n{stdout}\n"
+        f"Stderr:\n{stderr}"
+    )
+    metadata = {
+        "server_name": server_name,
+        "command": command,
+        "exit_code": exit_code,
+        "source": "command_history",
+    }
+    chunks = _chunk_text(doc)
+    for i, chunk in enumerate(chunks):
+        command_history.add(
+            ids=[f"cmd_{uuid.uuid4().hex[:12]}"],
+            documents=[chunk],
+            metadatas=[{**metadata, "chunk_index": i}],
+        )
+
+
+def index_log(server_name: str, log_source: str, content: str):
+    """Index server log output into server_logs collection."""
+    metadata = {
+        "server_name": server_name,
+        "log_source": log_source,
+        "source": "server_logs",
+    }
+    chunks = _chunk_text(content)
+    for i, chunk in enumerate(chunks):
+        server_logs.add(
+            ids=[f"log_{uuid.uuid4().hex[:12]}"],
+            documents=[chunk],
+            metadatas=[{**metadata, "chunk_index": i}],
+        )
+
+
+def index_config(server_name: str, file_path: str, content: str):
+    """Index a server config file into server_configs collection."""
+    metadata = {
+        "server_name": server_name,
+        "file_path": file_path,
+        "source": "server_configs",
+    }
+    chunks = _chunk_text(content)
+    for i, chunk in enumerate(chunks):
+        server_configs.add(
+            ids=[f"cfg_{uuid.uuid4().hex[:12]}"],
+            documents=[chunk],
+            metadatas=[{**metadata, "chunk_index": i}],
+        )
+
+
+# Map collection name -> collection object
+_COLLECTIONS = {
+    "command_history": command_history,
+    "server_logs": server_logs,
+    "server_configs": server_configs,
+}
+
+
+def search(
+    query: str,
+    collection_names: Optional[List[str]] = None,
+    n_results: int = 3,
+) -> str:
+    """
+    Semantic search across one or more knowledge collections.
+    Returns formatted string of top results for LLM consumption.
+    """
+    targets = collection_names or list(_COLLECTIONS.keys())
+    all_results = []
+
+    for name in targets:
+        col = _COLLECTIONS.get(name)
+        if not col or col.count() == 0:
+            continue
+        results = col.query(query_texts=[query], n_results=min(n_results, col.count()))
+        if results and results["documents"]:
+            for doc, meta in zip(results["documents"][0], results["metadatas"][0]):
+                all_results.append(
+                    f"[{meta.get('source', name)}] "
+                    f"Server: {meta.get('server_name', '?')} | "
+                    f"{doc[:300]}"
+                )
+
+    if not all_results:
+        return "No relevant knowledge found in the database."
+
+    return "\n---\n".join(all_results)

app/modules/servers/__init__.py

@@ -0,0 +1 @@
+from .models import Server