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/hosts.py

@@ -37,14 +37,43 @@ class HostRead(BaseModel):
 async def list_hosts(
     skip: int = 0,
     limit: int = 100,
+    hostname: Optional[str] = None,
+    ip_address: Optional[str] = None,
+    os: Optional[str] = None,
+    sort_by: Optional[str] = None,
+    sort_desc: bool = False,
     current_user: User = Depends(get_current_active_user),
     tenant_id: int = Depends(get_tenant_id),
     db: Session = Depends(get_db)
 ):
     """
-    List hosts scoped to user's tenant
+    List hosts scoped to user's tenant with advanced filtering
+    
+    Supports:
+    - Filtering by hostname, IP address, OS
+    - Sorting by any field
+    - Pagination
     """
-    hosts = db.query(Host).filter(Host.tenant_id == tenant_id).offset(skip).limit(limit).all()
+    query = db.query(Host).filter(Host.tenant_id == tenant_id)
+    
+    # Apply filters
+    if hostname:
+        query = query.filter(Host.hostname.ilike(f"%{hostname}%"))
+    if ip_address:
+        query = query.filter(Host.ip_address.ilike(f"%{ip_address}%"))
+    if os:
+        query = query.filter(Host.os.ilike(f"%{os}%"))
+    
+    # Apply sorting
+    if sort_by:
+        sort_column = getattr(Host, sort_by, None)
+        if sort_column:
+            if sort_desc:
+                query = query.order_by(sort_column.desc())
+            else:
+                query = query.order_by(sort_column)
+    
+    hosts = query.offset(skip).limit(limit).all()
     return hosts
 
 

backend/app/api/routes/notifications.py

@@ -0,0 +1,164 @@
+from typing import List
+from fastapi import APIRouter, Depends, HTTPException, status, WebSocket, WebSocketDisconnect
+from sqlalchemy.orm import Session
+import json
+
+from app.core.database import get_db
+from app.core.deps import get_current_active_user
+from app.models.user import User
+from app.models.notification import Notification
+from app.schemas.notification import NotificationRead, NotificationUpdate
+
+router = APIRouter()
+
+# Store active WebSocket connections
+active_connections: dict[int, List[WebSocket]] = {}
+
+
+@router.get("/", response_model=List[NotificationRead])
+async def list_notifications(
+    skip: int = 0,
+    limit: int = 50,
+    unread_only: bool = False,
+    current_user: User = Depends(get_current_active_user),
+    db: Session = Depends(get_db)
+):
+    """
+    List notifications for current user
+    
+    Supports filtering by read/unread status.
+    """
+    query = db.query(Notification).filter(
+        Notification.user_id == current_user.id,
+        Notification.tenant_id == current_user.tenant_id
+    )
+    
+    if unread_only:
+        query = query.filter(Notification.is_read == False)
+    
+    notifications = query.order_by(Notification.created_at.desc()).offset(skip).limit(limit).all()
+    return notifications
+
+
+@router.put("/{notification_id}", response_model=NotificationRead)
+async def update_notification(
+    notification_id: int,
+    notification_update: NotificationUpdate,
+    current_user: User = Depends(get_current_active_user),
+    db: Session = Depends(get_db)
+):
+    """
+    Update notification (mark as read/unread)
+    """
+    notification = db.query(Notification).filter(
+        Notification.id == notification_id,
+        Notification.user_id == current_user.id
+    ).first()
+    
+    if not notification:
+        raise HTTPException(
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail="Notification not found"
+        )
+    
+    notification.is_read = notification_update.is_read
+    db.commit()
+    db.refresh(notification)
+    
+    return notification
+
+
+@router.post("/mark-all-read")
+async def mark_all_read(
+    current_user: User = Depends(get_current_active_user),
+    db: Session = Depends(get_db)
+):
+    """
+    Mark all notifications as read for current user
+    """
+    db.query(Notification).filter(
+        Notification.user_id == current_user.id,
+        Notification.is_read == False
+    ).update({"is_read": True})
+    db.commit()
+    
+    return {"message": "All notifications marked as read"}
+
+
+@router.websocket("/ws")
+async def websocket_endpoint(
+    websocket: WebSocket,
+    db: Session = Depends(get_db)
+):
+    """
+    WebSocket endpoint for real-time notifications
+    
+    Clients should send their token on connect, then receive notifications in real-time.
+    """
+    await websocket.accept()
+    user_id = None
+    
+    try:
+        # Wait for authentication message
+        auth_data = await websocket.receive_text()
+        auth_json = json.loads(auth_data)
+        token = auth_json.get("token")
+        
+        # Validate token and get user (simplified - in production use proper auth)
+        from app.core.security import verify_token
+        payload = verify_token(token)
+        if payload:
+            user_id = payload.get("sub")
+            
+            # Register connection
+            if user_id not in active_connections:
+                active_connections[user_id] = []
+            active_connections[user_id].append(websocket)
+            
+            # Send confirmation
+            await websocket.send_json({"type": "connected", "user_id": user_id})
+            
+            # Keep connection alive
+            while True:
+                data = await websocket.receive_text()
+                # Echo back for keepalive
+                await websocket.send_json({"type": "pong"})
+        else:
+            await websocket.close(code=1008)  # Policy violation
+            
+    except WebSocketDisconnect:
+        # Remove connection
+        if user_id and user_id in active_connections:
+            active_connections[user_id].remove(websocket)
+            if not active_connections[user_id]:
+                del active_connections[user_id]
+    except Exception as e:
+        print(f"WebSocket error: {e}")
+        await websocket.close(code=1011)  # Internal error
+
+
+async def send_notification_to_user(user_id: int, notification: dict):
+    """
+    Send notification to all connected clients for a user
+    
+    Args:
+        user_id: User ID to send notification to
+        notification: Notification data to send
+    """
+    if user_id in active_connections:
+        disconnected = []
+        for websocket in active_connections[user_id]:
+            try:
+                await websocket.send_json({
+                    "type": "notification",
+                    "data": notification
+                })
+            except:
+                disconnected.append(websocket)
+        
+        # Clean up disconnected websockets
+        for ws in disconnected:
+            active_connections[user_id].remove(ws)
+        
+        if not active_connections[user_id]:
+            del active_connections[user_id]

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