215 lines
7.4 KiB
Python
215 lines
7.4 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from fastapi import FastAPI, HTTPException
|
|
|
|
from policy_distributor import build_policy_bundle
|
|
from policy_schema import PolicyActivateRequest, PolicyCreateRequest, PolicyStatusRequest, PolicyUpdateRequest
|
|
from policy_storage import PolicyStorage
|
|
|
|
|
|
def _env(name: str, default: str) -> str:
|
|
value = os.environ.get(name)
|
|
return value if value not in (None, "") else default
|
|
|
|
|
|
APP_NAME = "aw-dlp-policy-engine"
|
|
APP_VERSION = "0.1.0"
|
|
DB_PATH = _env("AW_DLP_POLICY_ENGINE_DB_PATH", "/var/lib/activitywatch/dlp-policy-engine.sqlite")
|
|
storage = PolicyStorage(DB_PATH)
|
|
|
|
app = FastAPI(title=APP_NAME, version=APP_VERSION)
|
|
AGENT_STATE: dict[str, dict[str, Any]] = {}
|
|
|
|
|
|
@app.get("/healthz")
|
|
def healthz() -> dict[str, str]:
|
|
return {
|
|
"status": "ok",
|
|
"service": APP_NAME,
|
|
"db_path": DB_PATH,
|
|
"db_exists": str(Path(DB_PATH).exists()).lower(),
|
|
}
|
|
|
|
|
|
@app.get("/api/0/dlp/policies")
|
|
def list_policies() -> dict[str, object]:
|
|
return {"items": storage.list_policies()}
|
|
|
|
|
|
@app.post("/api/0/dlp/policies", status_code=201)
|
|
def create_policy(payload: PolicyCreateRequest) -> dict[str, object]:
|
|
try:
|
|
item = storage.create_policy(
|
|
name=payload.name,
|
|
description=payload.description,
|
|
policy=payload.policy.model_dump(mode="json"),
|
|
activate=payload.activate,
|
|
actor=payload.actor,
|
|
)
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
return {"item": item}
|
|
|
|
|
|
@app.get("/api/0/dlp/policies/active")
|
|
def get_active_policy() -> dict[str, object]:
|
|
item = storage.get_active_policy()
|
|
if not item:
|
|
raise HTTPException(status_code=404, detail="no active policy configured")
|
|
return build_policy_bundle(item)
|
|
|
|
|
|
@app.get("/api/0/dlp/policies/active/version")
|
|
def get_active_policy_version() -> dict[str, object]:
|
|
item = storage.get_active_policy()
|
|
if not item:
|
|
raise HTTPException(status_code=404, detail="no active policy configured")
|
|
return {
|
|
"active": True,
|
|
"policyId": item["id"],
|
|
"version": item["current_version"],
|
|
"checksum": item["checksum"],
|
|
"updatedAtUtc": item["updated_at"],
|
|
}
|
|
|
|
|
|
@app.post("/api/0/dlp/policies/agents/{agent_id}/heartbeat")
|
|
def update_agent_policy_heartbeat(agent_id: str, payload: dict[str, Any]) -> dict[str, object]:
|
|
AGENT_STATE[agent_id] = {
|
|
"agentId": agent_id,
|
|
"hostname": payload.get("hostname") or agent_id,
|
|
"version": payload.get("version"),
|
|
"checksum": payload.get("checksum"),
|
|
"updatedAtUtc": payload.get("updatedAtUtc"),
|
|
}
|
|
return {"ok": True, "agent": AGENT_STATE[agent_id]}
|
|
|
|
|
|
@app.get("/api/0/dlp/policies/agents/{agent_id}/desired")
|
|
def get_agent_desired_policy(agent_id: str) -> dict[str, object]:
|
|
item = storage.get_active_policy()
|
|
if not item:
|
|
raise HTTPException(status_code=404, detail="no active policy configured")
|
|
|
|
current = AGENT_STATE.get(agent_id, {})
|
|
current_version = current.get("version")
|
|
current_checksum = current.get("checksum")
|
|
desired_version = item["current_version"]
|
|
desired_checksum = item["checksum"]
|
|
refresh_now = (str(current_version) != str(desired_version)) or (str(current_checksum) != str(desired_checksum))
|
|
|
|
return {
|
|
"agentId": agent_id,
|
|
"refreshNow": refresh_now,
|
|
"reason": "mismatch" if refresh_now else "up-to-date",
|
|
"current": {
|
|
"version": current_version,
|
|
"checksum": current_checksum,
|
|
},
|
|
"desired": {
|
|
"policyId": item["id"],
|
|
"version": desired_version,
|
|
"checksum": desired_checksum,
|
|
"updatedAtUtc": item["updated_at"],
|
|
},
|
|
}
|
|
|
|
|
|
@app.post("/api/0/dlp/policies/rollback")
|
|
def rollback_active_policy(payload: PolicyActivateRequest) -> dict[str, object]:
|
|
item = storage.rollback_active_policy(actor=payload.actor)
|
|
if not item:
|
|
raise HTTPException(status_code=404, detail="no active policy configured")
|
|
return {"item": item}
|
|
|
|
|
|
@app.get("/api/0/dlp/policies/{policy_id}")
|
|
def get_policy(policy_id: int) -> dict[str, object]:
|
|
item = storage.get_policy(policy_id)
|
|
if not item:
|
|
raise HTTPException(status_code=404, detail="policy not found")
|
|
return {"item": item}
|
|
|
|
|
|
@app.put("/api/0/dlp/policies/{policy_id}")
|
|
def update_policy(policy_id: int, payload: PolicyUpdateRequest) -> dict[str, object]:
|
|
try:
|
|
item = storage.update_policy(
|
|
policy_id=policy_id,
|
|
name=payload.name,
|
|
description=payload.description,
|
|
policy=payload.policy.model_dump(mode="json") if payload.policy is not None else None,
|
|
activate=payload.activate,
|
|
actor=payload.actor,
|
|
)
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
|
|
if not item:
|
|
raise HTTPException(status_code=404, detail="policy not found")
|
|
return {"item": item}
|
|
|
|
|
|
@app.post("/api/0/dlp/policies/{policy_id}/activate")
|
|
def activate_policy(policy_id: int, payload: PolicyActivateRequest) -> dict[str, object]:
|
|
try:
|
|
item = storage.activate_policy(policy_id=policy_id, actor=payload.actor)
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail=str(exc)) from exc
|
|
if not item:
|
|
raise HTTPException(status_code=404, detail="policy not found")
|
|
return {"item": item}
|
|
|
|
|
|
@app.post("/api/0/dlp/policies/{policy_id}/submit")
|
|
def submit_policy_for_approval(policy_id: int, payload: PolicyStatusRequest) -> dict[str, object]:
|
|
item = storage.set_policy_status(policy_id, "pending_approval", payload.actor, payload.comment)
|
|
if not item:
|
|
raise HTTPException(status_code=404, detail="policy not found")
|
|
return {"item": item}
|
|
|
|
|
|
@app.post("/api/0/dlp/policies/{policy_id}/approve")
|
|
def approve_policy(policy_id: int, payload: PolicyStatusRequest) -> dict[str, object]:
|
|
item = storage.set_policy_status(policy_id, "approved", payload.actor, payload.comment)
|
|
if not item:
|
|
raise HTTPException(status_code=404, detail="policy not found")
|
|
return {"item": item}
|
|
|
|
|
|
@app.post("/api/0/dlp/policies/{policy_id}/draft")
|
|
def return_policy_to_draft(policy_id: int, payload: PolicyStatusRequest) -> dict[str, object]:
|
|
item = storage.set_policy_status(policy_id, "draft", payload.actor, payload.comment)
|
|
if not item:
|
|
raise HTTPException(status_code=404, detail="policy not found")
|
|
return {"item": item}
|
|
|
|
|
|
@app.delete("/api/0/dlp/policies/{policy_id}")
|
|
def delete_policy(policy_id: int) -> dict[str, bool]:
|
|
try:
|
|
deleted = storage.delete_policy(policy_id)
|
|
except ValueError as exc:
|
|
raise HTTPException(status_code=409, detail=str(exc)) from exc
|
|
|
|
if not deleted:
|
|
raise HTTPException(status_code=404, detail="policy not found")
|
|
return {"deleted": True}
|
|
|
|
|
|
@app.get("/api/0/dlp/policies/audit")
|
|
def list_policy_audit(limit: int = 200) -> dict[str, object]:
|
|
return {"items": storage.list_audit(policy_id=None, limit=max(1, min(limit, 1000)))}
|
|
|
|
|
|
@app.get("/api/0/dlp/policies/{policy_id}/audit")
|
|
def list_single_policy_audit(policy_id: int, limit: int = 200) -> dict[str, object]:
|
|
if not storage.get_policy(policy_id):
|
|
raise HTTPException(status_code=404, detail="policy not found")
|
|
return {"items": storage.list_audit(policy_id=policy_id, limit=max(1, min(limit, 1000)))}
|