feat(agent): switch worktime sessions to rust agent

2026-06-04 10:36:19 +03:00
parent 152705a999
commit ff234af701
12 changed files with 455 additions and 9 deletions
@@ -294,6 +294,7 @@ dependencies = [
 "serde",
 "serde_json",
 "tempfile",
+ "windows-sys 0.59.0",
 ]

 [[package]]
@@ -2628,6 +2629,15 @@ dependencies = [
 "windows-targets 0.52.6",
 ]

+[[package]]
+name = "windows-sys"
+version = "0.59.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b"
+dependencies = [
+ "windows-targets 0.52.6",
+]
+
 [[package]]
 name = "windows-sys"
 version = "0.60.2"
@@ -14,6 +14,9 @@ reqwest.workspace = true
 serde.workspace = true
 serde_json.workspace = true

+[target.'cfg(windows)'.dependencies]
+windows-sys = { version = "0.59", features = ["Win32_System_RemoteDesktop"] }
+
 [dev-dependencies]
 tempfile.workspace = true

@@ -11,3 +11,8 @@ enable_workforce_activity = true
 spool_dir = "/var/lib/awatch-agent/spool"
 timeout_seconds = 10
 retry_attempts = 3
+
+# Optional ActivityWatch compatibility output.
+# Enables Rust replacement for the PowerShell worktime-session collector path.
+aw_api_base = "http://awatch.local:5600/api/0"
+aw_worktime_enabled = false
@@ -1,5 +1,6 @@
 use anyhow::Result;
 use chrono::Utc;
+use std::process::Command;

 use crate::collectors::common::{
    agent_id, command_output, current_session, domain, hostname, role_security_events, username,
@@ -37,16 +38,25 @@ impl TelemetryCollector for WindowsCollector {
    }

    fn collect_sessions(&self) -> Result<SessionSnapshot> {
-        let mut active = vec![current_session("local")];
-        let mut rdp = Vec::new();
+        let mut active = windows_query_user_sessions();
+        if active.is_empty() {
+            active.push(current_session("local"));
+        }
+        let mut rdp = active
+            .iter()
+            .filter(|session| session.session_type == "rdp")
+            .cloned()
+            .collect::<Vec<_>>();
        if std::env::var("SESSIONNAME")
            .unwrap_or_default()
            .to_ascii_lowercase()
            .contains("rdp")
        {
            let session = current_session("rdp");
-            rdp.push(session.clone());
-            active.push(session);
+            if !rdp.iter().any(|item| item.username == session.username) {
+                rdp.push(session.clone());
+                active.push(session);
+            }
        }
        Ok(SessionSnapshot {
            active_sessions: active,
@@ -105,6 +115,198 @@ fn windows_version() -> String {
        .unwrap_or_else(|| "Windows".to_string())
 }

+fn windows_query_user_sessions() -> Vec<crate::telemetry::SessionInfo> {
+    let native = windows_wts_sessions();
+    if !native.is_empty() {
+        return native;
+    }
+    let raw = command_output_utf16le("cmd", &["/U", "/C", "query user"])
+        .or_else(|| command_output_utf16le("cmd", &["/U", "/C", "quser"]))
+        .or_else(|| command_output_lossy_combined("cmd", &["/C", "query user"]))
+        .or_else(|| command_output_lossy_combined("cmd", &["/C", "quser"]))
+        .unwrap_or_default();
+    raw.lines()
+        .skip(1)
+        .filter_map(parse_query_user_line)
+        .collect()
+}
+
+#[cfg(windows)]
+fn windows_wts_sessions() -> Vec<crate::telemetry::SessionInfo> {
+    use std::ptr;
+    use windows_sys::Win32::System::RemoteDesktop::{
+        WTS_CURRENT_SERVER_HANDLE, WTS_SESSION_INFOW, WTSActive, WTSEnumerateSessionsW,
+        WTSFreeMemory, WTSUserName,
+    };
+
+    let mut sessions_ptr: *mut WTS_SESSION_INFOW = ptr::null_mut();
+    let mut count = 0_u32;
+    let ok = unsafe {
+        WTSEnumerateSessionsW(
+            WTS_CURRENT_SERVER_HANDLE,
+            0,
+            1,
+            &mut sessions_ptr,
+            &mut count,
+        )
+    };
+    if ok == 0 || sessions_ptr.is_null() || count == 0 {
+        return Vec::new();
+    }
+
+    let sessions =
+        unsafe { std::slice::from_raw_parts(sessions_ptr, usize::try_from(count).unwrap_or(0)) };
+    let mut items = Vec::new();
+    for session in sessions {
+        let username = wts_session_string(session.SessionId, WTSUserName);
+        if username.trim().is_empty() {
+            continue;
+        }
+        let station = unsafe { wide_nul_to_string(session.pWinStationName) };
+        let state = session.State;
+        let session_type = if station.to_ascii_lowercase().contains("rdp") {
+            "rdp"
+        } else {
+            "local"
+        };
+        items.push(crate::telemetry::SessionInfo {
+            session_id: session.SessionId.to_string(),
+            username,
+            session_type: session_type.to_string(),
+            remote_addr: None,
+            started_at: None,
+            active: state == WTSActive,
+        });
+    }
+    unsafe {
+        WTSFreeMemory(sessions_ptr.cast());
+    }
+    items
+}
+
+#[cfg(not(windows))]
+fn windows_wts_sessions() -> Vec<crate::telemetry::SessionInfo> {
+    Vec::new()
+}
+
+#[cfg(windows)]
+fn wts_session_string(session_id: u32, class: i32) -> String {
+    use std::ptr;
+    use windows_sys::Win32::System::RemoteDesktop::{
+        WTS_CURRENT_SERVER_HANDLE, WTSFreeMemory, WTSQuerySessionInformationW,
+    };
+
+    let mut buffer = ptr::null_mut();
+    let mut bytes = 0_u32;
+    let ok = unsafe {
+        WTSQuerySessionInformationW(
+            WTS_CURRENT_SERVER_HANDLE,
+            session_id,
+            class,
+            &mut buffer,
+            &mut bytes,
+        )
+    };
+    if ok == 0 || buffer.is_null() || bytes == 0 {
+        return String::new();
+    }
+    let len = usize::try_from(bytes / 2).unwrap_or(0);
+    let value = unsafe {
+        let slice = std::slice::from_raw_parts(buffer, len);
+        String::from_utf16_lossy(slice)
+            .trim_matches('\0')
+            .trim()
+            .to_string()
+    };
+    unsafe {
+        WTSFreeMemory(buffer.cast());
+    }
+    value
+}
+
+#[cfg(windows)]
+unsafe fn wide_nul_to_string(ptr: *const u16) -> String {
+    if ptr.is_null() {
+        return String::new();
+    }
+    let mut len = 0;
+    while unsafe { *ptr.add(len) } != 0 {
+        len += 1;
+    }
+    let slice = unsafe { std::slice::from_raw_parts(ptr, len) };
+    String::from_utf16_lossy(slice)
+}
+
+fn command_output_utf16le(program: &str, args: &[&str]) -> Option<String> {
+    let output = Command::new(program).args(args).output().ok()?;
+    if !output.status.success() {
+        return None;
+    }
+    let mut bytes = output.stdout;
+    bytes.extend_from_slice(&output.stderr);
+    if bytes.is_empty() {
+        return None;
+    }
+    let mut words = Vec::new();
+    for chunk in bytes.chunks_exact(2) {
+        words.push(u16::from_le_bytes([chunk[0], chunk[1]]));
+    }
+    String::from_utf16(&words)
+        .ok()
+        .map(|value| value.trim().to_string())
+        .filter(|value| !value.is_empty())
+}
+
+fn command_output_lossy_combined(program: &str, args: &[&str]) -> Option<String> {
+    let output = Command::new(program).args(args).output().ok()?;
+    if !output.status.success() {
+        return None;
+    }
+    let mut bytes = output.stdout;
+    bytes.extend_from_slice(&output.stderr);
+    Some(String::from_utf8_lossy(&bytes).trim().to_string()).filter(|value| !value.is_empty())
+}
+
+fn parse_query_user_line(line: &str) -> Option<crate::telemetry::SessionInfo> {
+    let cleaned = line.trim().trim_start_matches('>').trim();
+    if cleaned.is_empty() {
+        return None;
+    }
+    let parts = cleaned.split_whitespace().collect::<Vec<_>>();
+    if parts.len() < 3 {
+        return None;
+    }
+    let username = parts.first()?.to_string();
+    let (session_name, session_id, state) = if parts.get(1)?.chars().all(|ch| ch.is_ascii_digit()) {
+        ("".to_string(), *parts.get(1)?, *parts.get(2)?)
+    } else {
+        (
+            parts.get(1)?.to_string(),
+            *parts.get(2)?,
+            *parts.get(3).unwrap_or(&"Unknown"),
+        )
+    };
+    let active = session_state_active(state);
+    let session_type = if session_name.to_ascii_lowercase().contains("rdp") {
+        "rdp"
+    } else {
+        "local"
+    };
+    Some(crate::telemetry::SessionInfo {
+        session_id: session_id.to_string(),
+        username,
+        session_type: session_type.to_string(),
+        remote_addr: None,
+        started_at: None,
+        active,
+    })
+}
+
+fn session_state_active(state: &str) -> bool {
+    let lower = state.to_lowercase();
+    lower.contains("active") || lower.contains("актив")
+}
+
 fn windows_memory() -> (u64, u64) {
    let Some(raw) = command_output(
        "wmic",
@@ -222,3 +424,32 @@ fn split_host_port(value: &str) -> Option<(String, u16)> {
        port.parse().ok()?,
    ))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn parses_query_user_line_with_rdp_session() {
+        let session =
+            parse_query_user_line(" user1                 rdp-tcp#5           3  Active").unwrap();
+        assert_eq!(session.username, "user1");
+        assert_eq!(session.session_id, "3");
+        assert_eq!(session.session_type, "rdp");
+        assert!(session.active);
+    }
+
+    #[test]
+    fn parses_query_user_line_without_session_name() {
+        let session = parse_query_user_line(" user2                 4  Disc").unwrap();
+        assert_eq!(session.username, "user2");
+        assert_eq!(session.session_id, "4");
+        assert_eq!(session.session_type, "local");
+        assert!(!session.active);
+    }
+
+    #[test]
+    fn detects_russian_active_state() {
+        assert!(session_state_active("Активно"));
+    }
+}
@@ -16,6 +16,8 @@ pub struct AgentConfig {
    pub spool_dir: PathBuf,
    pub timeout_seconds: u64,
    pub retry_attempts: u32,
+    pub aw_api_base: Option<String>,
+    pub aw_worktime_enabled: bool,
 }

 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -49,6 +51,8 @@ impl Default for AgentConfig {
            spool_dir: default_spool_dir(),
            timeout_seconds: 10,
            retry_attempts: 3,
+            aw_api_base: None,
+            aw_worktime_enabled: false,
        }
    }
 }
@@ -103,6 +107,10 @@ impl AgentConfig {
                "spool_dir" => config.spool_dir = PathBuf::from(value),
                "timeout_seconds" => config.timeout_seconds = value.parse().unwrap_or(10),
                "retry_attempts" => config.retry_attempts = value.parse().unwrap_or(3),
+                "aw_api_base" => config.aw_api_base = Some(value.trim_end_matches('/').to_string()),
+                "aw_worktime_enabled" => {
+                    config.aw_worktime_enabled = parse_bool(value, false);
+                }
                _ => {}
            }
        }
@@ -131,6 +139,8 @@ api_key = "change-me"
 collect_interval_seconds = 30
 role = "firewall"
 enable_processes = false
+aw_api_base = "http://awatch.local:5600/api/0"
+aw_worktime_enabled = true
 spool_dir = "/tmp/awatch-spool"
 "#,
        )
@@ -138,6 +148,11 @@ spool_dir = "/tmp/awatch-spool"
        assert_eq!(config.role, AgentRole::Firewall);
        assert_eq!(config.collect_interval_seconds, 30);
        assert!(!config.enable_processes);
+        assert!(config.aw_worktime_enabled);
+        assert_eq!(
+            config.aw_api_base.as_deref(),
+            Some("http://awatch.local:5600/api/0")
+        );
        assert_eq!(config.spool_dir, PathBuf::from("/tmp/awatch-spool"));
    }
 }
@@ -10,7 +10,7 @@ use std::time::Duration;
 use anyhow::{Context, Result};
 use clap::Parser;
 use config::{AgentConfig, AgentRole, default_config_path};
-use transport::{TelemetryTransport, spool_health};
+use transport::{AwWorktimePublisher, TelemetryTransport, spool_health};

 #[derive(Debug, Parser)]
 #[command(about = "AWatch-rus Rust telemetry agent")]
@@ -27,6 +27,12 @@ struct Cli {
    #[arg(long, env = "AWATCH_AGENT_ROLE")]
    role: Option<String>,

+    #[arg(long, env = "AWATCH_AGENT_AW_API_BASE")]
+    aw_api_base: Option<String>,
+
+    #[arg(long, env = "AWATCH_AGENT_AW_WORKTIME_ENABLED")]
+    aw_worktime_enabled: Option<bool>,
+
    #[arg(long)]
    once: bool,

@@ -63,6 +69,12 @@ fn run() -> Result<i32> {
    if let Some(role) = cli.role {
        config.role = AgentRole::parse(&role);
    }
+    if let Some(aw_api_base) = cli.aw_api_base {
+        config.aw_api_base = Some(aw_api_base.trim_end_matches('/').to_string());
+    }
+    if let Some(enabled) = cli.aw_worktime_enabled {
+        config.aw_worktime_enabled = enabled;
+    }
    if cli.spool_health {
        println!(
            "{}",
@@ -72,6 +84,7 @@ fn run() -> Result<i32> {
    }

    let transport = TelemetryTransport::new(&config);
+    let aw_worktime = AwWorktimePublisher::new(&config);
    if cli.flush_spool {
        let flushed = transport.flush_spool()?;
        println!("{}", serde_json::json!({"ok": true, "flushed": flushed}));
@@ -86,6 +99,13 @@ fn run() -> Result<i32> {
        } else if let Err(err) = transport.send_or_spool(&record) {
            eprintln!("{err:#}");
        }
+        if !cli.print_json {
+            if let Some(publisher) = aw_worktime.as_ref() {
+                if let Err(err) = publisher.publish(&record) {
+                    eprintln!("{err:#}");
+                }
+            }
+        }
        if cli.once {
            break;
        }
@@ -9,7 +9,7 @@ use reqwest::blocking::Client;
 use reqwest::header::{HeaderMap, HeaderValue};

 use crate::config::AgentConfig;
-use crate::telemetry::TelemetryRecord;
+use crate::telemetry::{SessionInfo, TelemetryRecord};

 #[derive(Debug, Clone)]
 pub struct TelemetryTransport {
@@ -96,6 +96,147 @@ impl TelemetryTransport {
    }
 }

+#[derive(Debug, Clone)]
+pub struct AwWorktimePublisher {
+    aw_api_base: String,
+    timeout: Duration,
+}
+
+impl AwWorktimePublisher {
+    pub fn new(config: &AgentConfig) -> Option<Self> {
+        if !config.aw_worktime_enabled {
+            return None;
+        }
+        let aw_api_base = config
+            .aw_api_base
+            .as_ref()?
+            .trim_end_matches('/')
+            .to_string();
+        if aw_api_base.is_empty() {
+            return None;
+        }
+        Some(Self {
+            aw_api_base,
+            timeout: Duration::from_secs(config.timeout_seconds),
+        })
+    }
+
+    pub fn publish(&self, record: &TelemetryRecord) -> Result<usize> {
+        let client = Client::builder()
+            .timeout(self.timeout)
+            .build()
+            .context("build ActivityWatch HTTP client")?;
+        let bucket_id = format!(
+            "aw-worktime-sessions_{}",
+            sanitize_bucket_part(&record.hostname)
+        );
+        ensure_aw_bucket(
+            &client,
+            &self.aw_api_base,
+            &bucket_id,
+            "aw-worktime-session-collector",
+            "aw.worktime.session",
+            &record.hostname,
+        )?;
+        let sessions = if record.active_sessions.is_empty() {
+            vec![SessionInfo {
+                session_id: "0".to_string(),
+                username: record.username.clone(),
+                session_type: "local".to_string(),
+                remote_addr: None,
+                started_at: None,
+                active: true,
+            }]
+        } else {
+            record.active_sessions.clone()
+        };
+        let mut sent = 0;
+        let sample_seconds = 60_i64;
+        for session in sessions {
+            let payload = serde_json::json!({
+                "timestamp": record.timestamp,
+                "duration": sample_seconds,
+                "data": {
+                    "username": session.username,
+                    "userId": format!("{}\\{}", record.hostname, session.username),
+                    "sessionId": session_id_number(&session),
+                    "sessionName": session.session_type,
+                    "state": if session.active { "Active" } else { "Disconnected" },
+                    "active": session.active,
+                    "sampleSeconds": sample_seconds,
+                    "pollSeconds": sample_seconds,
+                    "hostname": record.hostname,
+                    "source": "awatch-agent-rs"
+                }
+            });
+            client
+                .post(format!(
+                    "{}/buckets/{}/heartbeat?pulsetime=180",
+                    self.aw_api_base, bucket_id
+                ))
+                .json(&payload)
+                .send()
+                .and_then(|response| response.error_for_status())
+                .context("publish ActivityWatch worktime heartbeat")?;
+            sent += 1;
+        }
+        Ok(sent)
+    }
+}
+
+fn ensure_aw_bucket(
+    client: &Client,
+    aw_api_base: &str,
+    bucket_id: &str,
+    client_name: &str,
+    bucket_type: &str,
+    hostname: &str,
+) -> Result<()> {
+    let bucket_url = format!("{}/buckets/{}", aw_api_base, bucket_id);
+    if client
+        .get(&bucket_url)
+        .send()
+        .and_then(|response| response.error_for_status())
+        .is_ok()
+    {
+        return Ok(());
+    }
+    let body = serde_json::json!({
+        "client": client_name,
+        "type": bucket_type,
+        "hostname": hostname,
+    });
+    client
+        .post(&bucket_url)
+        .json(&body)
+        .send()
+        .and_then(|response| response.error_for_status())
+        .context("create ActivityWatch worktime bucket")?;
+    Ok(())
+}
+
+fn session_id_number(session: &SessionInfo) -> i64 {
+    session
+        .session_id
+        .split(|ch: char| !ch.is_ascii_digit())
+        .find(|part| !part.is_empty())
+        .and_then(|part| part.parse::<i64>().ok())
+        .unwrap_or(0)
+}
+
+fn sanitize_bucket_part(value: &str) -> String {
+    value
+        .chars()
+        .map(|ch| {
+            if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' {
+                ch
+            } else {
+                '_'
+            }
+        })
+        .collect()
+}
+
 pub fn flush_spool_dir<F>(spool_dir: &Path, mut sender: F) -> Result<usize>
 where
    F: FnMut(&TelemetryRecord) -> Result<()>,
@@ -202,4 +343,17 @@ mod tests {
        assert_eq!(seen, 1);
        assert!(!path.exists());
    }
+
+    #[test]
+    fn session_id_number_extracts_numeric_id() {
+        let session = SessionInfo {
+            session_id: "rdp-12-user".to_string(),
+            username: "user".to_string(),
+            session_type: "rdp".to_string(),
+            remote_addr: None,
+            started_at: None,
+            active: true,
+        };
+        assert_eq!(session_id_number(&session), 12);
+    }
 }
@@ -995,6 +995,7 @@ function New-ActivityWatchDeploymentConfig {
            windowEnabled = $WindowEnabled
            fileOpsEnabled = $FileOpsEnabled
            emailEnabled = $false
+            worktimeSessionEnabled = $true
        }
        logging = [pscustomobject]@{
            localAgentLogsEnabled = $LocalAgentLogsEnabled
@@ -2033,7 +2034,10 @@ function Invoke-ActivityWatchRecoveryLoop {
                $sessionRecords = Get-ActivityWatchSessionRecords
                $stateRoot = [string]$config.paths.stateRoot
                $sessionCollectorScript = if ($config.paths.PSObject.Properties.Name -contains 'sessionCollectorScript') { [string]$config.paths.sessionCollectorScript } else { Join-Path $stateRoot 'worktime-session-collector.ps1' }
-                Start-ActivityWatchCollectorScriptGlobalIfNeeded -ScriptPath $sessionCollectorScript -ConfigPath $ConfigPath
+                $worktimeSessionEnabled = if ($config.PSObject.Properties.Name -contains 'collectors' -and $config.collectors.PSObject.Properties.Name -contains 'worktimeSessionEnabled') { [bool]$config.collectors.worktimeSessionEnabled } else { $true }
+                if ($worktimeSessionEnabled) {
+                    Start-ActivityWatchCollectorScriptGlobalIfNeeded -ScriptPath $sessionCollectorScript -ConfigPath $ConfigPath
+                }

                $configuredLiveTasksStarted = $false
                foreach ($taskDef in $taskDefs) {
@@ -488,9 +488,10 @@ function Invoke-GuardCycle {

    $worktimeAge = $bucketChecks["aw-worktime-sessions_$hostname"].ageSeconds
    $sessionCollectorScript = if ($config.paths.PSObject.Properties.Name -contains 'sessionCollectorScript') { [string]$config.paths.sessionCollectorScript } else { Join-Path $stateRoot 'worktime-session-collector.ps1' }
-    $sessionCollectorRunning = Test-ActivityWatchCollectorRunningGlobal -ScriptPath $sessionCollectorScript
+    $worktimeSessionEnabled = if ($config.PSObject.Properties.Name -contains 'collectors' -and $config.collectors.PSObject.Properties.Name -contains 'worktimeSessionEnabled') { [bool]$config.collectors.worktimeSessionEnabled } else { $true }
+    $sessionCollectorRunning = $worktimeSessionEnabled -and (Test-ActivityWatchCollectorRunningGlobal -ScriptPath $sessionCollectorScript)
    $headlessKey = 'headless:worktime-session'
-    $needsHeadlessAction = (-not $sessionCollectorRunning -or $null -eq $worktimeAge -or [int]$worktimeAge -gt $HeadlessMaxAgeSeconds)
+    $needsHeadlessAction = $worktimeSessionEnabled -and (-not $sessionCollectorRunning -or $null -eq $worktimeAge -or [int]$worktimeAge -gt $HeadlessMaxAgeSeconds)
    if ($needsHeadlessAction) {
        $key = $headlessKey
        $allowed = Test-ActionAllowed -Runtime $Runtime -Key $key -CooldownSeconds $ActionCooldownSeconds -WindowSeconds $RestartWindowSeconds -MaxCount $MaxRestarts
@@ -81,6 +81,7 @@ while ($true) {
        if ($collectors) {
            if ($collectors.PSObject.Properties.Name -contains 'fileOpsEnabled') { $startFileOps = [bool]$collectors.fileOpsEnabled }
            if ($collectors.PSObject.Properties.Name -contains 'emailEnabled')   { $startEmail   = [bool]$collectors.emailEnabled }
+            if ($collectors.PSObject.Properties.Name -contains 'worktimeSessionEnabled') { $startWorktime = [bool]$collectors.worktimeSessionEnabled }
        }
        if ($isSession0) {
            $startBrowser = $false
@@ -178,6 +178,7 @@ $report = [ordered]@{
        afkEnabled = $AfkEnabled
        windowEnabled = $WindowEnabled
        fileOpsEnabled = $FileOpsEnabled
+        worktimeSessionEnabled = $true
    }
    hardeningApplied = (-not $SkipHardening)
 }
@@ -94,6 +94,7 @@ $config = [pscustomobject]@{
        windowEnabled = $false
        fileOpsEnabled = $true
        emailEnabled = $true
+        worktimeSessionEnabled = $true
    }
    logging = [pscustomobject]@{
        localAgentLogsEnabled = $true