From 518c445157d6492260b78169343542a411c23c3c Mon Sep 17 00:00:00 2001 From: igor04091968 Date: Tue, 12 May 2026 01:01:08 +0300 Subject: [PATCH] feat(dlp-policy): add agent heartbeat/desired refresh channel for push-pull policy sync --- aw-server/dlp-policy-engine/policy_service.py | 58 +++++++++++++++ windows/dlp-endpoint-signals-collector.ps1 | 34 ++++++++- windows/dlp-policy-client.ps1 | 74 ++++++++++++++++++- 3 files changed, 163 insertions(+), 3 deletions(-) diff --git a/aw-server/dlp-policy-engine/policy_service.py b/aw-server/dlp-policy-engine/policy_service.py index 33840fc..555dd43 100644 --- a/aw-server/dlp-policy-engine/policy_service.py +++ b/aw-server/dlp-policy-engine/policy_service.py @@ -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) diff --git a/windows/dlp-endpoint-signals-collector.ps1 b/windows/dlp-endpoint-signals-collector.ps1 index 9aa0325..2d7921f 100644 --- a/windows/dlp-endpoint-signals-collector.ps1 +++ b/windows/dlp-endpoint-signals-collector.ps1 @@ -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() diff --git a/windows/dlp-policy-client.ps1 b/windows/dlp-policy-client.ps1 index 5b7f9f2..541df91 100644 --- a/windows/dlp-policy-client.ps1 +++ b/windows/dlp-policy-client.ps1 @@ -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