feat(dlp-policy): add agent heartbeat/desired refresh channel for push-pull policy sync

2026-05-12 01:01:08 +03:00
parent d812a8bb80
commit 518c445157
3 changed files with 163 additions and 3 deletions
@@ -3,6 +3,7 @@ from __future__ import annotations

 import os
 from pathlib import Path
+from typing import Any

 from fastapi import FastAPI, HTTPException

@@ -22,6 +23,7 @@ DB_PATH = _env("AW_DLP_POLICY_ENGINE_DB_PATH", "/var/lib/activitywatch/dlp-polic
 storage = PolicyStorage(DB_PATH)

 app = FastAPI(title=APP_NAME, version=APP_VERSION)
+AGENT_STATE: dict[str, dict[str, Any]] = {}


@app.get("/healthz")
@@ -62,6 +64,62 @@ def get_active_policy() -> dict[str, object]:
    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)
@@ -667,6 +667,29 @@ function Refresh-DlpPolicyFromServer {
    }
 }

+function Sync-DlpPolicyDesiredState {
+    if (-not $script:PolicyEngineEnabled -or -not $script:PolicyClientAvailable) {
+        return $false
+    }
+    if (-not $script:PolicyAgentId) {
+        return $false
+    }
+
+    try {
+        [void](Send-DlpPolicyAgentHeartbeat -ApiBase $script:PolicyApiBase -AgentId $script:PolicyAgentId -Hostname $script:Hostname -Version $script:PolicyVersion -Checksum $script:PolicyChecksum -TimeoutSec 10)
+        $desired = Get-RemoteDlpPolicyDesired -ApiBase $script:PolicyApiBase -AgentId $script:PolicyAgentId -TimeoutSec 10
+        if ($desired -and $desired.refreshNow -eq $true) {
+            Write-EndpointLog ("policy desired refresh requested: reason={0}" -f $desired.reason)
+            return (Refresh-DlpPolicyFromServer)
+        }
+        return $true
+    }
+    catch {
+        Write-EndpointLog ("policy desired sync failed: {0}" -f $_.Exception.Message)
+        return $false
+    }
+}
+
 function Initialize-DlpPolicy {
    if ($script:PolicyMode -eq 'server') {
        if (Refresh-DlpPolicyFromServer) {
@@ -1110,6 +1133,7 @@ $script:PolicyRefreshSeconds = [Math]::Max($resolvedPolicyRefreshSeconds, 60)
 $script:PolicyCachePath = $resolvedPolicyCachePath
 $script:LocalPolicyPath = $resolvedPolicyPath
 $script:LastPolicyRefreshAt = [datetime]::MinValue
+$script:PolicyAgentId = $resolvedHostname
 $script:TransportBackoffSeconds = 1
 # Integration test flag (backward compatible - defaults to false)
 $script:IntegrationTestEnabled = if ($deploymentConfig -and $deploymentConfig.PSObject.Properties.Name -contains 'integrationTestEnabled') { [bool]$deploymentConfig.integrationTestEnabled } else { $false }
@@ -1133,8 +1157,14 @@ while ($true) {
            Write-EndpointLog ("transport flush failed, backoff={0}s err={1}" -f $script:TransportBackoffSeconds, $_.Exception.Message)
        }

-        if ($script:PolicyMode -eq 'server' -and (($nowUtc = (Get-Date).ToUniversalTime()) - $script:LastPolicyRefreshAt).TotalSeconds -ge $script:PolicyRefreshSeconds) {
-            [void](Refresh-DlpPolicyFromServer)
+        if ($script:PolicyMode -eq 'server') {
+            $policyAge = ((Get-Date).ToUniversalTime() - $script:LastPolicyRefreshAt).TotalSeconds
+            if ($policyAge -ge $script:PolicyRefreshSeconds) {
+                [void](Refresh-DlpPolicyFromServer)
+            }
+            else {
+                [void](Sync-DlpPolicyDesiredState)
+            }
        }

        $nowUtc = (Get-Date).ToUniversalTime()
@@ -33,6 +33,49 @@ function Invoke-DlpPolicyGetJson {
    }
 }

+function Invoke-DlpPolicyPostJson {
+    param(
+        [Parameter(Mandatory = $true)][string]$Uri,
+        [Parameter(Mandatory = $true)]$Body,
+        [int]$TimeoutSec = 10
+    )
+
+    $json = $Body | ConvertTo-Json -Depth 10 -Compress
+    $bytes = [System.Text.Encoding]::UTF8.GetBytes($json)
+    $request = [System.Net.HttpWebRequest]::Create($Uri)
+    $request.Method = 'POST'
+    $request.Accept = 'application/json'
+    $request.ContentType = 'application/json'
+    $request.KeepAlive = $false
+    $request.Timeout = $TimeoutSec * 1000
+    $request.ReadWriteTimeout = $TimeoutSec * 1000
+    $request.ContentLength = $bytes.Length
+
+    $stream = $request.GetRequestStream()
+    try {
+        $stream.Write($bytes, 0, $bytes.Length)
+    }
+    finally {
+        $stream.Close()
+    }
+
+    $response = $request.GetResponse()
+    try {
+        $reader = New-Object System.IO.StreamReader($response.GetResponseStream(), [System.Text.Encoding]::UTF8)
+        try {
+            $raw = $reader.ReadToEnd()
+            if ($raw) { return ($raw | ConvertFrom-Json) }
+            return $null
+        }
+        finally {
+            $reader.Close()
+        }
+    }
+    finally {
+        $response.Close()
+    }
+}
+
 function Get-RemoteDlpPolicyBundle {
    param(
        [Parameter(Mandatory = $true)][string]$ApiBase,
@@ -52,6 +95,35 @@ function Get-RemoteDlpPolicyBundle {
    return $bundle
 }

+function Get-RemoteDlpPolicyDesired {
+    param(
+        [Parameter(Mandatory = $true)][string]$ApiBase,
+        [Parameter(Mandatory = $true)][string]$AgentId,
+        [int]$TimeoutSec = 10
+    )
+    $path = '/dlp/policies/agents/{0}/desired' -f [uri]::EscapeDataString($AgentId)
+    return Invoke-DlpPolicyGetJson -Uri ($ApiBase.TrimEnd('/') + $path) -TimeoutSec $TimeoutSec
+}
+
+function Send-DlpPolicyAgentHeartbeat {
+    param(
+        [Parameter(Mandatory = $true)][string]$ApiBase,
+        [Parameter(Mandatory = $true)][string]$AgentId,
+        [string]$Hostname,
+        [string]$Version,
+        [string]$Checksum,
+        [int]$TimeoutSec = 10
+    )
+    $path = '/dlp/policies/agents/{0}/heartbeat' -f [uri]::EscapeDataString($AgentId)
+    $body = @{
+        hostname = $Hostname
+        version = $Version
+        checksum = $Checksum
+        updatedAtUtc = (Get-Date).ToUniversalTime().ToString('o')
+    }
+    return Invoke-DlpPolicyPostJson -Uri ($ApiBase.TrimEnd('/') + $path) -Body $body -TimeoutSec $TimeoutSec
+}
+
 function Read-CachedDlpPolicyBundle {
    param([Parameter(Mandatory = $true)][string]$CachePath)

@@ -82,4 +154,4 @@ function Save-CachedDlpPolicyBundle {
    Set-Content -LiteralPath $CachePath -Value $json -Encoding UTF8
 }

-Export-ModuleMember -Function Invoke-DlpPolicyGetJson, Get-RemoteDlpPolicyBundle, Read-CachedDlpPolicyBundle, Save-CachedDlpPolicyBundle
+Export-ModuleMember -Function Invoke-DlpPolicyGetJson, Get-RemoteDlpPolicyBundle, Get-RemoteDlpPolicyDesired, Send-DlpPolicyAgentHeartbeat, Read-CachedDlpPolicyBundle, Save-CachedDlpPolicyBundle