diff --git a/backend/app/api/v1/router.py b/backend/app/api/v1/router.py index fc561b4..b6e79e6 100644 --- a/backend/app/api/v1/router.py +++ b/backend/app/api/v1/router.py @@ -4,9 +4,9 @@ import io from ipaddress import ip_address, ip_interface, ip_network from pathlib import Path -from fastapi import APIRouter, Depends, Header, HTTPException, Request +from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request from fastapi.responses import FileResponse, PlainTextResponse, StreamingResponse -from sqlalchemy import func, select +from sqlalchemy import func, or_, select from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session @@ -744,29 +744,24 @@ async def policy_deployment_status(db: Session, policy: Policy) -> dict[str, obj def dashboard_top_talkers(db: Session) -> list[dict[str, int | str]]: totals: dict[str, int] = {} - workloads = db.scalars(select(Workload)).all() - workloads_by_node_vmid = {(workload.node_id, workload.external_id): workload for workload in workloads} - ip_owners = { - address.address: db.get(Workload, address.workload_id) - for address in db.scalars(select(IpAddress).where(IpAddress.workload_id.is_not(None))).all() - } - - flows = db.scalars(select(TrafficFlow)).all() - for flow in flows: - raw = flow.raw if isinstance(flow.raw, dict) else {} - candidates: list[Workload] = [] - raw_vmid = str(raw.get("vmid") or "") - if raw_vmid: - workload = workloads_by_node_vmid.get((flow.node_id, raw_vmid)) - if workload: - candidates.append(workload) - for owner in (ip_owners.get(flow.source_ip), ip_owners.get(flow.destination_ip)): - if owner and owner not in candidates: - candidates.append(owner) - for workload in candidates: - totals[workload.name] = totals.get(workload.name, 0) + int(flow.bytes or 0) + source_rows = db.execute( + select(Workload.name, func.coalesce(func.sum(TrafficFlow.bytes), 0)) + .join(IpAddress, IpAddress.workload_id == Workload.id) + .join(TrafficFlow, TrafficFlow.source_ip == IpAddress.address) + .group_by(Workload.name) + ).all() + destination_rows = db.execute( + select(Workload.name, func.coalesce(func.sum(TrafficFlow.bytes), 0)) + .join(IpAddress, IpAddress.workload_id == Workload.id) + .join(TrafficFlow, TrafficFlow.destination_ip == IpAddress.address) + .group_by(Workload.name) + ).all() + for name, bytes_value in [*source_rows, *destination_rows]: + totals[str(name)] = totals.get(str(name), 0) + int(bytes_value or 0) if not totals: + workloads = db.scalars(select(Workload)).all() + workloads_by_node_vmid = {(workload.node_id, workload.external_id): workload for workload in workloads} agents = db.scalars(select(NodeAgent)).all() for agent in agents: payload = agent.last_payload if isinstance(agent.last_payload, dict) else {} @@ -799,8 +794,19 @@ def dashboard_suspicious_traffic(db: Session) -> list[dict[str, int | str]]: address.address for address in db.scalars(select(IpAddress).where(IpAddress.workload_id.is_not(None))).all() } + blocked_states = ["blocked", "drop", "dropped", "reject", "rejected", "deny", "denied"] + candidate_flows = db.scalars( + select(TrafficFlow) + .where( + or_( + TrafficFlow.destination_port.in_(list(sensitive_ports)), + func.lower(func.coalesce(TrafficFlow.state, "")).in_(blocked_states), + ) + ) + .order_by(TrafficFlow.updated_at.desc()) + ).all() events: dict[tuple[str, str, int], dict[str, int | str]] = {} - for flow in db.scalars(select(TrafficFlow).order_by(TrafficFlow.updated_at.desc())).all(): + for flow in candidate_flows: port = flow.destination_port or 0 source_internal = bool(subnet_label_for_ip(subnets, flow.source_ip)) destination_internal = bool(subnet_label_for_ip(subnets, flow.destination_ip)) or flow.destination_ip in workload_ips @@ -1507,7 +1513,12 @@ def workloads(_: CurrentUser, db: Session = Depends(get_db)) -> list[Workload]: @api_router.get("/vms/{workload_id}/insights", response_model=WorkloadInsight) -async def workload_insights(workload_id: str, _: CurrentUser, db: Session = Depends(get_db)) -> WorkloadInsight: +async def workload_insights( + workload_id: str, + _: CurrentUser, + db: Session = Depends(get_db), + traffic: str = Query(default="summary", pattern="^(summary|full)$"), +) -> WorkloadInsight: workload = db.get(Workload, workload_id) if not workload: raise HTTPException(status_code=404, detail="Workload not found") @@ -1517,10 +1528,9 @@ async def workload_insights(workload_id: str, _: CurrentUser, db: Session = Depe assigned_ips = db.scalars(select(IpAddress).where(IpAddress.workload_id == workload.id).order_by(IpAddress.address)).all() workload_ips = [address.address for address in assigned_ips] all_assigned_ips = db.scalars(select(IpAddress).where(IpAddress.workload_id.is_not(None))).all() - ip_owners = { - address.address: db.get(Workload, address.workload_id) - for address in all_assigned_ips - } + owner_ids = {address.workload_id for address in all_assigned_ips if address.workload_id} + owners = {workload.id: workload for workload in db.scalars(select(Workload).where(Workload.id.in_(owner_ids))).all()} if owner_ids else {} + ip_owners = {address.address: owners.get(address.workload_id) for address in all_assigned_ips if address.workload_id} workload_ips_by_id: dict[str, set[str]] = {} for address in all_assigned_ips: if address.workload_id: @@ -1530,11 +1540,14 @@ async def workload_insights(workload_id: str, _: CurrentUser, db: Session = Depe traffic = [] if workload_ips: workload_ip_set = set(workload_ips) - flows = db.scalars( + flow_query = ( select(TrafficFlow) .where((TrafficFlow.source_ip.in_(workload_ips)) | (TrafficFlow.destination_ip.in_(workload_ips))) .order_by((TrafficFlow.state == "blocked").desc(), TrafficFlow.updated_at.desc()) - ).all() + ) + if traffic == "summary": + flow_query = flow_query.limit(250) + flows = db.scalars(flow_query).all() for flow in flows: source_owner = ip_owners.get(flow.source_ip) destination_owner = ip_owners.get(flow.destination_ip) diff --git a/backend/app/main.py b/backend/app/main.py index 2a47687..a80bf71 100644 --- a/backend/app/main.py +++ b/backend/app/main.py @@ -1,5 +1,6 @@ from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware +from sqlalchemy import text from app.api.v1.router import api_router from app.core.config import get_settings @@ -7,6 +8,20 @@ from app.db.session import Base, SessionLocal, engine from app.seed.demo import seed_demo_data +def ensure_runtime_indexes() -> None: + index_statements = [ + "CREATE INDEX IF NOT EXISTS ix_ip_addresses_address ON ip_addresses (address)", + "CREATE INDEX IF NOT EXISTS ix_traffic_flows_node_updated ON traffic_flows (node_id, updated_at)", + "CREATE INDEX IF NOT EXISTS ix_traffic_flows_source_updated ON traffic_flows (source_ip, updated_at)", + "CREATE INDEX IF NOT EXISTS ix_traffic_flows_destination_updated ON traffic_flows (destination_ip, updated_at)", + "CREATE INDEX IF NOT EXISTS ix_traffic_flows_destination_port_updated ON traffic_flows (destination_port, updated_at)", + "CREATE INDEX IF NOT EXISTS ix_traffic_flows_state_updated ON traffic_flows (state, updated_at)", + ] + with engine.begin() as connection: + for statement in index_statements: + connection.execute(text(statement)) + + def create_app() -> FastAPI: settings = get_settings() app = FastAPI( @@ -27,6 +42,7 @@ def create_app() -> FastAPI: @app.on_event("startup") def startup() -> None: Base.metadata.create_all(bind=engine) + ensure_runtime_indexes() with SessionLocal() as db: seed_demo_data(db) @@ -39,4 +55,3 @@ def create_app() -> FastAPI: app = create_app() - diff --git a/backend/app/models/domain.py b/backend/app/models/domain.py index bcbea48..4d3bb1a 100644 --- a/backend/app/models/domain.py +++ b/backend/app/models/domain.py @@ -2,7 +2,7 @@ from datetime import datetime from enum import StrEnum from uuid import uuid4 -from sqlalchemy import JSON, BigInteger, Boolean, DateTime, Enum, ForeignKey, Integer, String, Text, UniqueConstraint +from sqlalchemy import JSON, BigInteger, Boolean, DateTime, Enum, ForeignKey, Index, Integer, String, Text, UniqueConstraint from sqlalchemy.orm import Mapped, mapped_column, relationship from app.db.session import Base @@ -185,7 +185,10 @@ class Subnet(Base, TimestampMixin): class IpAddress(Base, TimestampMixin): __tablename__ = "ip_addresses" - __table_args__ = (UniqueConstraint("subnet_id", "address"),) + __table_args__ = ( + UniqueConstraint("subnet_id", "address"), + Index("ix_ip_addresses_address", "address"), + ) id: Mapped[str] = mapped_column(String, primary_key=True, default=new_id) subnet_id: Mapped[str] = mapped_column(ForeignKey("subnets.id"), index=True) @@ -197,6 +200,13 @@ class IpAddress(Base, TimestampMixin): class TrafficFlow(Base, TimestampMixin): __tablename__ = "traffic_flows" + __table_args__ = ( + Index("ix_traffic_flows_node_updated", "node_id", "updated_at"), + Index("ix_traffic_flows_source_updated", "source_ip", "updated_at"), + Index("ix_traffic_flows_destination_updated", "destination_ip", "updated_at"), + Index("ix_traffic_flows_destination_port_updated", "destination_port", "updated_at"), + Index("ix_traffic_flows_state_updated", "state", "updated_at"), + ) id: Mapped[str] = mapped_column(String, primary_key=True, default=new_id) node_id: Mapped[str] = mapped_column(ForeignKey("nodes.id"), index=True) diff --git a/frontend/src/pages/Workloads.tsx b/frontend/src/pages/Workloads.tsx index f658e1a..8bb584f 100644 --- a/frontend/src/pages/Workloads.tsx +++ b/frontend/src/pages/Workloads.tsx @@ -452,8 +452,8 @@ export function Workloads() { const [selectedId, setSelectedId] = useState(""); const selected = selectedId || workloads.data?.[0]?.id || ""; const insight = useQuery({ - queryKey: ["workload-insight", selected], - queryFn: () => api(`/vms/${selected}/insights`), + queryKey: ["workload-insight", selected, "summary"], + queryFn: () => api(`/vms/${selected}/insights?traffic=summary`), enabled: Boolean(selected), }); const traffic = useMemo(() => summarizeTraffic(insight.data?.traffic ?? []), [insight.data?.traffic]); @@ -540,8 +540,8 @@ export function Workloads() { export function WorkloadDetail() { const { workloadId } = useParams(); const insight = useQuery({ - queryKey: ["workload-insight", workloadId], - queryFn: () => api(`/vms/${workloadId}/insights`), + queryKey: ["workload-insight", workloadId, "summary"], + queryFn: () => api(`/vms/${workloadId}/insights?traffic=summary`), enabled: Boolean(workloadId), }); const traffic = useMemo(() => summarizeTraffic(insight.data?.traffic ?? []), [insight.data?.traffic]); @@ -638,8 +638,8 @@ export function WorkloadFlows() { const [port, setPort] = useState(""); const [page, setPage] = useState(1); const insight = useQuery({ - queryKey: ["workload-insight", workloadId], - queryFn: () => api(`/vms/${workloadId}/insights`), + queryKey: ["workload-insight", workloadId, "full"], + queryFn: () => api(`/vms/${workloadId}/insights?traffic=full`), enabled: Boolean(workloadId), }); const traffic = useMemo(() => summarizeTraffic(insight.data?.traffic ?? []), [insight.data?.traffic]);