feat: add traffic flow deduplication with 24h retention and update-in-place for existing flows
Add traffic_flow_key helper to generate unique flow identifier from node/IPs/protocol/ports/decision, implement 24-hour retention cutoff to delete old flows instead of all flows on heartbeat, build existing_flows lookup map from database with composite key matching, update agent_heartbeat to check for existing flows and update bytes/packets/state/observed_at/raw in-place instead of creating duplicates, extend
This commit is contained in:
@@ -536,6 +536,19 @@ def raw_flow_decision(flow: TrafficFlow) -> str | None:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def traffic_flow_key(node_id: str, raw_flow: dict[str, object]) -> tuple[object, ...]:
|
||||||
|
decision = str(raw_flow.get("decision") or raw_flow.get("state") or "observed").lower()
|
||||||
|
return (
|
||||||
|
node_id,
|
||||||
|
str(raw_flow.get("source_ip") or ""),
|
||||||
|
str(raw_flow.get("destination_ip") or ""),
|
||||||
|
str(raw_flow.get("protocol") or "unknown"),
|
||||||
|
flow_int(raw_flow.get("source_port"), 0) or None,
|
||||||
|
flow_int(raw_flow.get("destination_port"), 0) or None,
|
||||||
|
decision,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
def policy_read_payload(policy: Policy, deployment_status: dict[str, object] | None = None) -> dict[str, object]:
|
def policy_read_payload(policy: Policy, deployment_status: dict[str, object] | None = None) -> dict[str, object]:
|
||||||
return {
|
return {
|
||||||
"id": policy.id,
|
"id": policy.id,
|
||||||
@@ -1357,29 +1370,53 @@ def agent_heartbeat(payload: AgentHeartbeat, authorization: str | None = Header(
|
|||||||
agent.version = payload.version
|
agent.version = payload.version
|
||||||
agent.last_seen_at = datetime.utcnow()
|
agent.last_seen_at = datetime.utcnow()
|
||||||
agent.last_payload = payload.model_dump(mode="json")
|
agent.last_payload = payload.model_dump(mode="json")
|
||||||
for old_flow in db.scalars(select(TrafficFlow).where(TrafficFlow.node_id == node.id)).all():
|
retention_cutoff = datetime.utcnow() - timedelta(hours=24)
|
||||||
|
for old_flow in db.scalars(select(TrafficFlow).where(TrafficFlow.node_id == node.id, TrafficFlow.updated_at < retention_cutoff)).all():
|
||||||
db.delete(old_flow)
|
db.delete(old_flow)
|
||||||
|
existing_flows = {
|
||||||
|
(
|
||||||
|
flow.node_id,
|
||||||
|
flow.source_ip,
|
||||||
|
flow.destination_ip,
|
||||||
|
flow.protocol,
|
||||||
|
flow.source_port,
|
||||||
|
flow.destination_port,
|
||||||
|
str(flow.state or "observed").lower(),
|
||||||
|
): flow
|
||||||
|
for flow in db.scalars(select(TrafficFlow).where(TrafficFlow.node_id == node.id)).all()
|
||||||
|
}
|
||||||
for raw_flow in payload.flows[:1000]:
|
for raw_flow in payload.flows[:1000]:
|
||||||
source_ip = str(raw_flow.get("source_ip") or "")
|
source_ip = str(raw_flow.get("source_ip") or "")
|
||||||
destination_ip = str(raw_flow.get("destination_ip") or "")
|
destination_ip = str(raw_flow.get("destination_ip") or "")
|
||||||
if not source_ip or not destination_ip:
|
if not source_ip or not destination_ip:
|
||||||
continue
|
continue
|
||||||
observed_at = payload.collected_at or datetime.utcnow()
|
observed_at = payload.collected_at or datetime.utcnow()
|
||||||
db.add(
|
observed_at = observed_at.replace(tzinfo=None) if observed_at.tzinfo else observed_at
|
||||||
TrafficFlow(
|
key = traffic_flow_key(node.id, raw_flow)
|
||||||
node_id=node.id,
|
existing = existing_flows.get(key)
|
||||||
source_ip=source_ip,
|
if existing:
|
||||||
destination_ip=destination_ip,
|
existing.bytes = flow_int(raw_flow.get("bytes"))
|
||||||
protocol=str(raw_flow.get("protocol") or "unknown"),
|
existing.packets = flow_int(raw_flow.get("packets"))
|
||||||
source_port=flow_int(raw_flow.get("source_port"), 0) or None,
|
existing.state = str(raw_flow.get("decision") or raw_flow.get("state") or "") or None
|
||||||
destination_port=flow_int(raw_flow.get("destination_port"), 0) or None,
|
existing.observed_at = observed_at
|
||||||
bytes=flow_int(raw_flow.get("bytes")),
|
existing.raw = raw_flow
|
||||||
packets=flow_int(raw_flow.get("packets")),
|
existing.updated_at = datetime.utcnow()
|
||||||
state=str(raw_flow.get("decision") or raw_flow.get("state") or "") or None,
|
continue
|
||||||
observed_at=observed_at.replace(tzinfo=None) if observed_at.tzinfo else observed_at,
|
flow = TrafficFlow(
|
||||||
raw=raw_flow,
|
node_id=node.id,
|
||||||
)
|
source_ip=source_ip,
|
||||||
|
destination_ip=destination_ip,
|
||||||
|
protocol=str(raw_flow.get("protocol") or "unknown"),
|
||||||
|
source_port=flow_int(raw_flow.get("source_port"), 0) or None,
|
||||||
|
destination_port=flow_int(raw_flow.get("destination_port"), 0) or None,
|
||||||
|
bytes=flow_int(raw_flow.get("bytes")),
|
||||||
|
packets=flow_int(raw_flow.get("packets")),
|
||||||
|
state=str(raw_flow.get("decision") or raw_flow.get("state") or "") or None,
|
||||||
|
observed_at=observed_at,
|
||||||
|
raw=raw_flow,
|
||||||
)
|
)
|
||||||
|
db.add(flow)
|
||||||
|
existing_flows[key] = flow
|
||||||
commit_or_400(db)
|
commit_or_400(db)
|
||||||
return {"status": "ok", "node_id": node.id}
|
return {"status": "ok", "node_id": node.id}
|
||||||
|
|
||||||
@@ -1416,7 +1453,7 @@ async def workload_insights(workload_id: str, _: CurrentUser, db: Session = Depe
|
|||||||
flows = db.scalars(
|
flows = db.scalars(
|
||||||
select(TrafficFlow)
|
select(TrafficFlow)
|
||||||
.where((TrafficFlow.source_ip.in_(workload_ips)) | (TrafficFlow.destination_ip.in_(workload_ips)))
|
.where((TrafficFlow.source_ip.in_(workload_ips)) | (TrafficFlow.destination_ip.in_(workload_ips)))
|
||||||
.order_by(TrafficFlow.updated_at.desc())
|
.order_by((TrafficFlow.state == "blocked").desc(), TrafficFlow.updated_at.desc())
|
||||||
.limit(50)
|
.limit(50)
|
||||||
).all()
|
).all()
|
||||||
for flow in flows:
|
for flow in flows:
|
||||||
|
|||||||
Reference in New Issue
Block a user