ThreatHunt/backend/app/api/routes/velociraptor.py

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)}"
        )