awatch-rus/AWatch-rus
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(agent): harden rust worktime fallback mode

Browse Source
This commit is contained in:
igor04091968
2026-06-04 11:21:54 +03:00
parent ff234af701
commit d7616fd9a9
13 changed files with 491 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
adk-rust/crates/awatch-agent-rs/src/collectors/common.rs
+1
View File
@@ -63,6 +63,7 @@ pub fn current_session(session_type: &str) -> SessionInfo {
session_id: format!("{}-{}", session_type, username()),
username: username(),
session_type: session_type.to_string(),
session_source: Some("env_sessionname_fallback".to_string()),
remote_addr: std::env::var("SSH_CLIENT")
.ok()
.and_then(|value| value.split_whitespace().next().map(str::to_string)),
adk-rust/crates/awatch-agent-rs/src/collectors/freebsd.rs
+6 -1
View File
@@ -10,7 +10,7 @@ use crate::config::AgentRole;
use crate::telemetry::{
IdentityInfo, NetworkConnectionInfo, NetworkInterfaceInfo, NetworkSnapshot, ProcessInfo,
ResourceInfo, SecurityEventInfo, SessionSnapshot, TelemetryCollector, WorkforceActivityInfo,
empty_workforce_activity,
dedupe_sessions, diagnostics_for_sessions, empty_workforce_activity,
};
#[derive(Debug, Clone)]
@@ -46,10 +46,15 @@ impl TelemetryCollector for FreeBsdCollector {
ssh.push(session.clone());
active.push(session);
}
let host = hostname();
let active = dedupe_sessions(&host, active);
let ssh = dedupe_sessions(&host, ssh);
let diagnostics = diagnostics_for_sessions(&active, &[], "env_sessionname_fallback", None);
Ok(SessionSnapshot {
active_sessions: active,
rdp_sessions: Vec::new(),
ssh_sessions: ssh,
diagnostics,
})
}
adk-rust/crates/awatch-agent-rs/src/collectors/linux.rs
+6 -1
View File
@@ -12,7 +12,7 @@ use crate::config::AgentRole;
use crate::telemetry::{
IdentityInfo, NetworkConnectionInfo, NetworkInterfaceInfo, NetworkSnapshot, ProcessInfo,
ResourceInfo, SecurityEventInfo, SessionSnapshot, TelemetryCollector, WorkforceActivityInfo,
empty_workforce_activity,
dedupe_sessions, diagnostics_for_sessions, empty_workforce_activity,
};
#[derive(Debug, Clone)]
@@ -49,10 +49,15 @@ impl TelemetryCollector for LinuxCollector {
ssh.push(session.clone());
active.push(session);
}
let host = hostname();
let active = dedupe_sessions(&host, active);
let ssh = dedupe_sessions(&host, ssh);
let diagnostics = diagnostics_for_sessions(&active, &[], "env_sessionname_fallback", None);
Ok(SessionSnapshot {
active_sessions: active,
rdp_sessions: Vec::new(),
ssh_sessions: ssh,
diagnostics,
})
}
adk-rust/crates/awatch-agent-rs/src/collectors/windows.rs
+101 -20
View File
@@ -8,8 +8,8 @@ use crate::collectors::common::{
use crate::config::AgentRole;
use crate::telemetry::{
IdentityInfo, NetworkConnectionInfo, NetworkSnapshot, ProcessInfo, ResourceInfo,
SecurityEventInfo, SessionSnapshot, TelemetryCollector, WorkforceActivityInfo,
empty_workforce_activity,
SecurityEventInfo, SessionSnapshot, TelemetryCollector, WorkforceActivityInfo, dedupe_sessions,
diagnostics_for_sessions, empty_workforce_activity,
};
#[derive(Debug, Clone)]
@@ -38,10 +38,27 @@ impl TelemetryCollector for WindowsCollector {
}
fn collect_sessions(&self) -> Result<SessionSnapshot> {
let mut active = windows_query_user_sessions();
if active.is_empty() {
active.push(current_session("local"));
let host = hostname();
let mut collection = windows_query_user_sessions();
if collection.sessions.is_empty()
&& std::env::var("SESSIONNAME")
.unwrap_or_default()
.to_ascii_lowercase()
.contains("rdp")
{
let mut session = current_session("rdp");
session.session_source = Some("env_sessionname_fallback".to_string());
collection.sessions.push(session);
collection.source = "env_sessionname_fallback".to_string();
}
if collection.sessions.is_empty() {
let mut session = current_session("local");
session.session_source = Some("local_fallback".to_string());
collection.sessions.push(session);
collection.source = "local_fallback".to_string();
collection.error = Some("WTS API and quser did not return sessions".to_string());
}
let mut active = dedupe_sessions(&host, collection.sessions);
let mut rdp = active
.iter()
.filter(|session| session.session_type == "rdp")
@@ -52,16 +69,26 @@ impl TelemetryCollector for WindowsCollector {
.to_ascii_lowercase()
.contains("rdp")
{
let session = current_session("rdp");
if !rdp.iter().any(|item| item.username == session.username) {
rdp.push(session.clone());
active.push(session);
}
let mut merged = active.clone();
merged.push(with_session_source(
current_session("rdp"),
"env_sessionname_fallback",
));
active = dedupe_sessions(&host, merged);
rdp = active
.iter()
.filter(|session| session.session_type == "rdp")
.cloned()
.collect::<Vec<_>>();
}
rdp = dedupe_sessions(&host, rdp);
let diagnostics =
diagnostics_for_sessions(&active, &rdp, collection.source, collection.error);
Ok(SessionSnapshot {
active_sessions: active,
rdp_sessions: rdp,
ssh_sessions: Vec::new(),
diagnostics,
})
}
@@ -115,19 +142,65 @@ fn windows_version() -> String {
.unwrap_or_else(|| "Windows".to_string())
}
fn windows_query_user_sessions() -> Vec<crate::telemetry::SessionInfo> {
#[derive(Debug)]
struct SessionCollection {
sessions: Vec<crate::telemetry::SessionInfo>,
source: String,
error: Option<String>,
}
fn with_session_source(
mut session: crate::telemetry::SessionInfo,
source: &str,
) -> crate::telemetry::SessionInfo {
session.session_source = Some(source.to_string());
session
}
fn windows_query_user_sessions() -> SessionCollection {
let native = windows_wts_sessions();
if !native.is_empty() {
return native;
return SessionCollection {
sessions: native,
source: "wts_api".to_string(),
error: None,
};
}
let raw = command_output_utf16le("cmd", &["/U", "/C", "query user"])
if let Some(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"]))
{
let sessions = parse_query_user_sessions(&raw, "quser_utf16");
if !sessions.is_empty() {
return SessionCollection {
sessions,
source: "quser_utf16".to_string(),
error: None,
};
}
}
if let Some(raw) = command_output_lossy_combined("cmd", &["/C", "query user"])
.or_else(|| command_output_lossy_combined("cmd", &["/C", "quser"]))
.unwrap_or_default();
{
let sessions = parse_query_user_sessions(&raw, "quser_lossy");
if !sessions.is_empty() {
return SessionCollection {
sessions,
source: "quser_lossy".to_string(),
error: None,
};
}
}
SessionCollection {
sessions: Vec::new(),
source: "local_fallback".to_string(),
error: Some("WTS API and quser returned no sessions".to_string()),
}
}
fn parse_query_user_sessions(raw: &str, source: &str) -> Vec<crate::telemetry::SessionInfo> {
raw.lines()
.skip(1)
.filter_map(parse_query_user_line)
.filter_map(|line| parse_query_user_line(line, source))
.collect()
}
@@ -173,6 +246,7 @@ fn windows_wts_sessions() -> Vec<crate::telemetry::SessionInfo> {
session_id: session.SessionId.to_string(),
username,
session_type: session_type.to_string(),
session_source: Some("wts_api".to_string()),
remote_addr: None,
started_at: None,
active: state == WTSActive,
@@ -267,7 +341,7 @@ fn command_output_lossy_combined(program: &str, args: &[&str]) -> Option<String>
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> {
fn parse_query_user_line(line: &str, source: &str) -> Option<crate::telemetry::SessionInfo> {
let cleaned = line.trim().trim_start_matches('>').trim();
if cleaned.is_empty() {
return None;
@@ -296,6 +370,7 @@ fn parse_query_user_line(line: &str) -> Option<crate::telemetry::SessionInfo> {
session_id: session_id.to_string(),
username,
session_type: session_type.to_string(),
session_source: Some(source.to_string()),
remote_addr: None,
started_at: None,
active,
@@ -431,20 +506,26 @@ mod tests {
#[test]
fn parses_query_user_line_with_rdp_session() {
let session =
parse_query_user_line(" user1 rdp-tcp#5 3 Active").unwrap();
let session = parse_query_user_line(
" user1 rdp-tcp#5 3 Active",
"quser_utf16",
)
.unwrap();
assert_eq!(session.username, "user1");
assert_eq!(session.session_id, "3");
assert_eq!(session.session_type, "rdp");
assert_eq!(session.session_source.as_deref(), Some("quser_utf16"));
assert!(session.active);
}
#[test]
fn parses_query_user_line_without_session_name() {
let session = parse_query_user_line(" user2 4 Disc").unwrap();
let session =
parse_query_user_line(" user2 4 Disc", "quser_lossy").unwrap();
assert_eq!(session.username, "user2");
assert_eq!(session.session_id, "4");
assert_eq!(session.session_type, "local");
assert_eq!(session.session_source.as_deref(), Some("quser_lossy"));
assert!(!session.active);
}
adk-rust/crates/awatch-agent-rs/src/main.rs
+1 -1
View File
@@ -101,7 +101,7 @@ fn run() -> Result<i32> {
}
if !cli.print_json {
if let Some(publisher) = aw_worktime.as_ref() {
if let Err(err) = publisher.publish(&record) {
if let Err(err) = publisher.publish_or_spool(&record) {
eprintln!("{err:#}");
}
}
adk-rust/crates/awatch-agent-rs/src/telemetry.rs
+64
View File
@@ -1,6 +1,7 @@
use anyhow::Result;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::BTreeSet;
pub const COLLECTOR_VERSION: &str = env!("CARGO_PKG_VERSION");
@@ -26,6 +27,7 @@ pub struct TelemetryRecord {
pub network_connections: Vec<NetworkConnectionInfo>,
pub workforce_activity: WorkforceActivityInfo,
pub security_events: Vec<SecurityEventInfo>,
pub diagnostics: AgentDiagnostics,
pub collector_version: String,
}
@@ -53,11 +55,23 @@ pub struct SessionInfo {
pub session_id: String,
pub username: String,
pub session_type: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub session_source: Option<String>,
pub remote_addr: Option<String>,
pub started_at: Option<DateTime<Utc>>,
pub active: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct AgentDiagnostics {
pub sessions_collected_total: usize,
pub rdp_sessions_total: usize,
pub active_sessions_total: usize,
pub collector_source: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub collector_error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ProcessInfo {
pub pid: u32,
@@ -117,6 +131,7 @@ pub struct SessionSnapshot {
pub active_sessions: Vec<SessionInfo>,
pub rdp_sessions: Vec<SessionInfo>,
pub ssh_sessions: Vec<SessionInfo>,
pub diagnostics: AgentDiagnostics,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
@@ -160,11 +175,43 @@ pub trait TelemetryCollector {
network_connections: network.connections,
workforce_activity: self.collect_workforce_activity()?,
security_events: self.collect_security_events()?,
diagnostics: sessions.diagnostics,
collector_version: COLLECTOR_VERSION.to_string(),
})
}
}
pub fn diagnostics_for_sessions(
active_sessions: &[SessionInfo],
rdp_sessions: &[SessionInfo],
collector_source: impl Into<String>,
collector_error: Option<String>,
) -> AgentDiagnostics {
AgentDiagnostics {
sessions_collected_total: active_sessions.len(),
rdp_sessions_total: rdp_sessions.len(),
active_sessions_total: active_sessions
.iter()
.filter(|session| session.active)
.count(),
collector_source: collector_source.into(),
collector_error,
}
}
pub fn dedupe_sessions(hostname: &str, sessions: Vec<SessionInfo>) -> Vec<SessionInfo> {
let mut seen = BTreeSet::new();
sessions
.into_iter()
.filter(|session| {
seen.insert(format!(
"{}\u{1f}{}\u{1f}{}\u{1f}{}",
hostname, session.username, session.session_id, session.session_type
))
})
.collect()
}
pub fn empty_workforce_activity() -> WorkforceActivityInfo {
WorkforceActivityInfo {
active_today: false,
@@ -204,11 +251,28 @@ mod tests {
network_connections: Vec::new(),
workforce_activity: empty_workforce_activity(),
security_events: Vec::new(),
diagnostics: diagnostics_for_sessions(&[], &[], "test", None),
collector_version: COLLECTOR_VERSION.to_string(),
};
let value = serde_json::to_value(record).unwrap();
assert_eq!(value["agent_id"], "agent-1");
assert!(value.get("network_connections").unwrap().is_array());
assert!(value.get("workforce_activity").is_some());
assert_eq!(value["diagnostics"]["collector_source"], "test");
}
#[test]
fn deduplicates_sessions_by_host_user_id_and_type() {
let session = SessionInfo {
session_id: "2".to_string(),
username: "user".to_string(),
session_type: "rdp".to_string(),
session_source: Some("wts_api".to_string()),
remote_addr: None,
started_at: None,
active: true,
};
let deduped = dedupe_sessions("HOST-EXAMPLE", vec![session.clone(), session]);
assert_eq!(deduped.len(), 1);
}
}
adk-rust/crates/awatch-agent-rs/src/transport.rs
+106 -14
View File
@@ -99,7 +99,9 @@ impl TelemetryTransport {
#[derive(Debug, Clone)]
pub struct AwWorktimePublisher {
aw_api_base: String,
spool_dir: PathBuf,
timeout: Duration,
retry_attempts: u32,
}
impl AwWorktimePublisher {
@@ -117,10 +119,25 @@ impl AwWorktimePublisher {
}
Some(Self {
aw_api_base,
spool_dir: config.spool_dir.join("aw-worktime"),
timeout: Duration::from_secs(config.timeout_seconds),
retry_attempts: config.retry_attempts,
})
}
pub fn publish_or_spool(&self, record: &TelemetryRecord) -> Result<()> {
if let Err(err) = self.flush_spool() {
eprintln!("ActivityWatch worktime spool flush failed: {err:#}");
}
match self.publish(record) {
Ok(_) => Ok(()),
Err(err) => {
self.spool(record)?;
Err(err)
}
}
}
pub fn publish(&self, record: &TelemetryRecord) -> Result<usize> {
let client = Client::builder()
.timeout(self.timeout)
@@ -143,6 +160,7 @@ impl AwWorktimePublisher {
session_id: "0".to_string(),
username: record.username.clone(),
session_type: "local".to_string(),
session_source: Some("local_fallback".to_string()),
remote_addr: None,
started_at: None,
active: true,
@@ -161,27 +179,52 @@ impl AwWorktimePublisher {
"userId": format!("{}\\{}", record.hostname, session.username),
"sessionId": session_id_number(&session),
"sessionName": session.session_type,
"sessionSource": session.session_source,
"state": if session.active { "Active" } else { "Disconnected" },
"active": session.active,
"sampleSeconds": sample_seconds,
"pollSeconds": sample_seconds,
"hostname": record.hostname,
"source": "awatch-agent-rs"
"source": "awatch-agent-rs",
"collectorSource": record.diagnostics.collector_source,
"sessionsCollectedTotal": record.diagnostics.sessions_collected_total,
"rdpSessionsTotal": record.diagnostics.rdp_sessions_total,
"activeSessionsTotal": record.diagnostics.active_sessions_total,
"collectorError": record.diagnostics.collector_error,
}
});
client
.post(format!(
post_json_with_retry(
&client,
&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")?;
),
&payload,
self.retry_attempts,
)
.context("publish ActivityWatch worktime heartbeat")?;
sent += 1;
}
Ok(sent)
}
pub fn spool(&self, record: &TelemetryRecord) -> Result<PathBuf> {
fs::create_dir_all(&self.spool_dir)
.with_context(|| format!("create worktime spool {}", self.spool_dir.display()))?;
let file_name = format!(
"{}-{}.json",
record.timestamp.format("%Y%m%dT%H%M%S%.3fZ"),
sanitize_file_part(&record.agent_id)
);
let path = self.spool_dir.join(file_name);
fs::write(&path, serde_json::to_vec(record)?)
.with_context(|| format!("write worktime spool {}", path.display()))?;
Ok(path)
}
pub fn flush_spool(&self) -> Result<usize> {
flush_spool_dir(&self.spool_dir, |record| self.publish(record).map(|_| ()))
}
}
fn ensure_aw_bucket(
@@ -206,15 +249,42 @@ fn ensure_aw_bucket(
"type": bucket_type,
"hostname": hostname,
});
client
.post(&bucket_url)
.json(&body)
.send()
.and_then(|response| response.error_for_status())
post_json_with_retry(client, &bucket_url, &body, 3)
.context("create ActivityWatch worktime bucket")?;
Ok(())
}
fn post_json_with_retry(
client: &Client,
url: &str,
payload: &serde_json::Value,
retry_attempts: u32,
) -> Result<()> {
let mut last_error = None;
for attempt in 0..retry_attempts.max(1) {
let result = client
.post(url)
.json(payload)
.send()
.and_then(|response| response.error_for_status())
.map(|_| ());
match result {
Ok(()) => return Ok(()),
Err(err) => {
last_error = Some(err);
let backoff = Duration::from_millis(250 * u64::from(attempt + 1));
thread::sleep(backoff);
}
}
}
Err(anyhow!(
"HTTP POST failed: {}",
last_error
.map(|err| err.to_string())
.unwrap_or_else(|| "unknown error".to_string())
))
}
fn session_id_number(session: &SessionInfo) -> i64 {
session
.session_id
@@ -295,7 +365,7 @@ mod tests {
use tempfile::tempdir;
use super::*;
use crate::telemetry::{TelemetryRecord, empty_workforce_activity};
use crate::telemetry::{TelemetryRecord, diagnostics_for_sessions, empty_workforce_activity};
fn record() -> TelemetryRecord {
TelemetryRecord {
@@ -319,6 +389,7 @@ mod tests {
network_connections: Vec::new(),
workforce_activity: empty_workforce_activity(),
security_events: Vec::new(),
diagnostics: diagnostics_for_sessions(&[], &[], "test", None),
collector_version: "test".to_string(),
}
}
@@ -350,10 +421,31 @@ mod tests {
session_id: "rdp-12-user".to_string(),
username: "user".to_string(),
session_type: "rdp".to_string(),
session_source: Some("wts_api".to_string()),
remote_addr: None,
started_at: None,
active: true,
};
assert_eq!(session_id_number(&session), 12);
}
#[test]
fn worktime_publisher_is_disabled_by_default() {
assert!(AwWorktimePublisher::new(&AgentConfig::default()).is_none());
}
#[test]
fn worktime_publisher_spools_to_separate_dir() {
let dir = tempdir().unwrap();
let config = AgentConfig {
aw_api_base: Some("http://127.0.0.1:9/api/0".to_string()),
aw_worktime_enabled: true,
spool_dir: dir.path().to_path_buf(),
..AgentConfig::default()
};
let publisher = AwWorktimePublisher::new(&config).unwrap();
let path = publisher.spool(&record()).unwrap();
assert!(path.starts_with(dir.path().join("aw-worktime")));
assert!(path.is_file());
}
}