Implement Phase 3: Advanced search, real-time notifications, and Velociraptor integration

Co-authored-by: mblanke <9078342+mblanke@users.noreply.github.com>

backend/app/api/routes/velociraptor.py

@@ -0,0 +1,197 @@
+from typing import List, Dict, Any, Optional
+from fastapi import APIRouter, Depends, HTTPException, status
+from sqlalchemy.orm import Session
+from pydantic import BaseModel
+
+from app.core.database import get_db
+from app.core.deps import get_current_active_user, require_role
+from app.core.velociraptor import get_velociraptor_client
+from app.models.user import User
+
+router = APIRouter()
+
+
+class VelociraptorConfig(BaseModel):
+    """Velociraptor server configuration"""
+    base_url: str
+    api_key: str
+
+
+class ArtifactCollectionRequest(BaseModel):
+    """Request to collect an artifact"""
+    client_id: str
+    artifact_name: str
+    parameters: Optional[Dict[str, Any]] = None
+
+
+class HuntCreateRequest(BaseModel):
+    """Request to create a hunt"""
+    hunt_name: str
+    artifact_name: str
+    description: str
+    parameters: Optional[Dict[str, Any]] = None
+
+
+# In a real implementation, this would be stored per-tenant in database
+# For now, using a simple in-memory store
+velociraptor_configs: Dict[int, VelociraptorConfig] = {}
+
+
+@router.post("/config")
+async def set_velociraptor_config(
+    config: VelociraptorConfig,
+    current_user: User = Depends(require_role(["admin"])),
+    db: Session = Depends(get_db)
+):
+    """
+    Configure Velociraptor integration (admin only)
+    
+    Stores Velociraptor server URL and API key for the tenant.
+    """
+    velociraptor_configs[current_user.tenant_id] = config
+    return {"message": "Velociraptor configuration saved"}
+
+
+@router.get("/clients")
+async def list_velociraptor_clients(
+    limit: int = 100,
+    current_user: User = Depends(get_current_active_user),
+    db: Session = Depends(get_db)
+):
+    """
+    List clients from Velociraptor server
+    """
+    config = velociraptor_configs.get(current_user.tenant_id)
+    if not config:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Velociraptor not configured for this tenant"
+        )
+    
+    client = get_velociraptor_client(config.base_url, config.api_key)
+    try:
+        clients = await client.list_clients(limit=limit)
+        return {"clients": clients}
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to fetch clients: {str(e)}"
+        )
+
+
+@router.get("/clients/{client_id}")
+async def get_velociraptor_client_info(
+    client_id: str,
+    current_user: User = Depends(get_current_active_user),
+    db: Session = Depends(get_db)
+):
+    """
+    Get information about a specific Velociraptor client
+    """
+    config = velociraptor_configs.get(current_user.tenant_id)
+    if not config:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Velociraptor not configured for this tenant"
+        )
+    
+    client = get_velociraptor_client(config.base_url, config.api_key)
+    try:
+        client_info = await client.get_client(client_id)
+        return client_info
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to fetch client: {str(e)}"
+        )
+
+
+@router.post("/collect")
+async def collect_artifact(
+    request: ArtifactCollectionRequest,
+    current_user: User = Depends(get_current_active_user),
+    db: Session = Depends(get_db)
+):
+    """
+    Collect an artifact from a Velociraptor client
+    """
+    config = velociraptor_configs.get(current_user.tenant_id)
+    if not config:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Velociraptor not configured for this tenant"
+        )
+    
+    client = get_velociraptor_client(config.base_url, config.api_key)
+    try:
+        result = await client.collect_artifact(
+            request.client_id,
+            request.artifact_name,
+            request.parameters
+        )
+        return result
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to collect artifact: {str(e)}"
+        )
+
+
+@router.post("/hunts")
+async def create_hunt(
+    request: HuntCreateRequest,
+    current_user: User = Depends(require_role(["admin"])),
+    db: Session = Depends(get_db)
+):
+    """
+    Create a new hunt (admin only)
+    """
+    config = velociraptor_configs.get(current_user.tenant_id)
+    if not config:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Velociraptor not configured for this tenant"
+        )
+    
+    client = get_velociraptor_client(config.base_url, config.api_key)
+    try:
+        result = await client.create_hunt(
+            request.hunt_name,
+            request.artifact_name,
+            request.description,
+            request.parameters
+        )
+        return result
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to create hunt: {str(e)}"
+        )
+
+
+@router.get("/hunts/{hunt_id}/results")
+async def get_hunt_results(
+    hunt_id: str,
+    limit: int = 1000,
+    current_user: User = Depends(get_current_active_user),
+    db: Session = Depends(get_db)
+):
+    """
+    Get results from a hunt
+    """
+    config = velociraptor_configs.get(current_user.tenant_id)
+    if not config:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Velociraptor not configured for this tenant"
+        )
+    
+    client = get_velociraptor_client(config.base_url, config.api_key)
+    try:
+        results = await client.get_hunt_results(hunt_id, limit=limit)
+        return {"results": results}
+    except Exception as e:
+        raise HTTPException(
+            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+            detail=f"Failed to fetch hunt results: {str(e)}"
+        )