Files
AWatch-rus/adk-rust/crates/check-aw-full/src/main.rs
T

609 lines
17 KiB
Rust

use std::net::{TcpStream, ToSocketAddrs};
use std::time::Duration;
use anyhow::{Context, Result};
use chrono::{DateTime, Utc};
use clap::Parser;
use reqwest::blocking::Client;
use serde_json::Value;
const DEFAULT_SERVER: &str = "http://10.10.10.13:5600";
const DEFAULT_HOST: &str = "SHARKON2025";
const DEFAULT_RDP_HOST: &str = "192.168.100.19";
const BUCKETS: &[(&str, &str)] = &[
("aw-watcher-afk", "AFK watcher"),
("aw-watcher-window", "Window watcher"),
("aw-worktime-sessions", "Worktime sessions"),
("aw-session-events", "Session events"),
("aw-dlp-endpoint-signals", "DLP signals"),
("aw-dlp-incidents", "DLP incidents"),
("aw-dlp-review", "DLP review"),
("aw-dlp-rules", "DLP rules"),
];
#[derive(Debug, Parser)]
#[command(about = "Full read-only ActivityWatch check for server, buckets, and RDP host")]
struct Cli {
#[arg(long, default_value = DEFAULT_SERVER)]
server: String,
#[arg(long, default_value = DEFAULT_HOST)]
host: String,
#[arg(long, default_value = DEFAULT_RDP_HOST)]
rdp_host: String,
#[arg(long, default_value_t = 15)]
timeout_seconds: u64,
#[arg(long)]
no_color: bool,
#[arg(long, default_value_t = true)]
dlp_enabled: bool,
}
#[derive(Debug, Clone)]
struct Colors {
red: &'static str,
green: &'static str,
yellow: &'static str,
cyan: &'static str,
reset: &'static str,
}
impl Colors {
fn new(enabled: bool) -> Self {
if enabled {
Self {
red: "\u{1b}[0;31m",
green: "\u{1b}[0;32m",
yellow: "\u{1b}[1;33m",
cyan: "\u{1b}[0;36m",
reset: "\u{1b}[0m",
}
} else {
Self {
red: "",
green: "",
yellow: "",
cyan: "",
reset: "",
}
}
}
fn paint(&self, color: &str, text: &str) -> String {
format!("{color}{text}{}", self.reset)
}
}
#[derive(Debug, Default)]
struct ContextState {
host_inactive: bool,
guard_healthy: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum BucketStatus {
Fresh,
Stale,
Dead,
Empty,
EventDriven,
Inactive,
Unknown,
}
#[derive(Debug, Clone)]
struct BucketRow {
label: &'static str,
last_id: String,
age: String,
status: BucketStatus,
}
#[derive(Debug, Default)]
struct Summary {
fresh: usize,
stale: usize,
dead: usize,
}
fn main() {
let code = match run() {
Ok(code) => code,
Err(err) => {
eprintln!("{err:#}");
1
}
};
std::process::exit(code);
}
fn run() -> Result<i32> {
let mut cli = Cli::parse();
if cli.server == DEFAULT_SERVER {
if let Some(value) = env_nonempty("CHECK_AW_FULL_SERVER")
.or_else(|| env_nonempty("AW_SMOKE_AW_SERVER"))
.or_else(|| env_nonempty("AW_SERVER"))
{
cli.server = value;
}
}
if cli.host == DEFAULT_HOST {
if let Some(value) = env_nonempty("CHECK_AW_FULL_HOST")
.or_else(|| env_nonempty("AW_SMOKE_SOURCE_HOSTNAME"))
.or_else(|| env_nonempty("AW_LOGICAL_HOST_ID"))
.or_else(|| env_nonempty("AW_MONITORED_WINDOWS_HOSTNAME"))
{
cli.host = value;
}
}
if cli.rdp_host == DEFAULT_RDP_HOST {
if let Some(value) = env_nonempty("CHECK_AW_FULL_RDP_HOST")
.or_else(|| env_nonempty("AW_SMOKE_WINDOWS_HOST"))
.or_else(|| env_nonempty("AW_WINDOWS_HOST"))
{
cli.rdp_host = value;
}
}
if let Some(value) =
env_nonempty("AW_DLP_ENABLED").or_else(|| env_nonempty("DETMIR_DLP_ENABLED"))
{
cli.dlp_enabled = parse_env_flag(&value);
}
let server = cli.server.trim_end_matches('/').to_string();
let colors = Colors::new(!cli.no_color && std::env::var_os("NO_COLOR").is_none());
let timeout = Duration::from_secs(cli.timeout_seconds.max(1));
let client = Client::builder()
.timeout(timeout)
.no_proxy()
.pool_max_idle_per_host(0)
.build()
.context("build HTTP client")?;
let now = Utc::now();
println!(
"{}",
colors.paint(
colors.cyan,
&format!("=== ActivityWatch Full Check: {} ===", cli.host)
)
);
println!();
println!(
"{}",
colors.paint(colors.cyan, &format!("--- 1. AW Server ({server}) ---"))
);
print!(" Connectivity... ");
let info = match get_json(&client, &format!("{server}/api/0/info")) {
Ok(value) => value,
Err(_) => {
println!(" {}", colors.paint(colors.red, "FAILED"));
return Ok(1);
}
};
let Some(version) = info.get("version").and_then(Value::as_str) else {
println!(" {}", colors.paint(colors.red, "FAILED"));
return Ok(1);
};
println!(
" {} (aw-server {version})",
colors.paint(colors.green, "OK")
);
print!(" CORS... ");
let cors_code = check_cors(&client, &server);
if cors_code == 200 {
println!("{}", colors.paint(colors.green, "OK"));
} else {
println!("{} (HTTP {cors_code})", colors.paint(colors.red, "FAIL"));
}
println!();
let context = read_context(&client, &server, &cli.host, now);
println!("{}", colors.paint(colors.cyan, "--- 2. Data Buckets ---"));
println!(
" {:<42} {:<8} {:<20} STATUS",
"BUCKET", "EVENTS", "LAST EVENT"
);
println!(
" {:<42} {:<8} {:<20} ------",
"------------------------------------------", "--------", "--------------------"
);
let mut rows = Vec::new();
for (bucket, label) in BUCKETS
.iter()
.copied()
.filter(|(bucket, _)| cli.dlp_enabled || !bucket.starts_with("aw-dlp-"))
{
let row = read_bucket_row(&client, &server, &cli.host, bucket, label, now, &context);
println!(
" {:<42} {:<8} {:<20} {}",
row.label,
row.last_id,
row.age,
render_status(&colors, row.status)
);
rows.push(row);
}
if !cli.dlp_enabled {
println!(
" {:<42} {:<8} {:<20} {}",
"DLP buckets",
"-",
"disabled",
colors.paint(colors.cyan, "SKIPPED")
);
}
println!();
println!(
"{}",
colors.paint(
colors.cyan,
&format!("--- 3. RDP Host ({}) ---", cli.rdp_host)
)
);
print!(" WinRM (5985)... ");
if tcp_open(&cli.rdp_host, 5985, Duration::from_secs(5)) {
println!("{}", colors.paint(colors.green, "OK"));
} else {
println!("{}", colors.paint(colors.red, "UNREACHABLE"));
}
print!(" SSH (22)... ");
if tcp_open(&cli.rdp_host, 22, Duration::from_secs(5)) {
println!("{}", colors.paint(colors.green, "OK"));
} else {
println!(
"{} (normal for Windows)",
colors.paint(colors.yellow, "CLOSED")
);
}
println!();
let summary = summarize(&rows);
println!("{}", colors.paint(colors.cyan, "--- 4. Summary ---"));
println!(
" FRESH: {}",
colors.paint(colors.green, &summary.fresh.to_string())
);
println!(
" STALE: {}",
colors.paint(colors.yellow, &summary.stale.to_string())
);
println!(
" DEAD: {}",
colors.paint(colors.red, &summary.dead.to_string())
);
if summary.dead > 0 || summary.stale > 0 {
println!();
println!(
" {} Some collectors may need restart on RDP host",
colors.paint(colors.red, "WARNING:")
);
println!(
" Run: {}",
colors.paint(
colors.cyan,
"ansible -i ansible/inventory.ini rdp-prod -m win_shell -a 'schtasks /Run /TN \"ActivityWatch Recovery\"'"
)
);
}
println!();
println!("{}", colors.paint(colors.cyan, "=== Check Complete ==="));
println!(" Timestamp: {}", now.format("%Y-%m-%d %H:%M:%S UTC"));
Ok(0)
}
fn read_context(client: &Client, server: &str, host: &str, now: DateTime<Utc>) -> ContextState {
let mut state = ContextState::default();
if let Ok(Some(event)) = latest_event(client, server, &format!("aw-worktime-sessions_{host}")) {
if let Some(ts) = event_timestamp(&event) {
let age = (now - ts).num_seconds();
let active = event
.pointer("/data/active")
.and_then(Value::as_bool)
.unwrap_or(false);
if (0..900).contains(&age) && !active {
state.host_inactive = true;
}
}
}
if let Ok(Some(event)) = latest_event(client, server, &format!("aw-rus-collector-guard_{host}"))
{
if let Some(ts) = event_timestamp(&event) {
let age = (now - ts).num_seconds();
let status = event
.pointer("/data/status")
.and_then(Value::as_str)
.unwrap_or("");
let problems = event
.pointer("/data/problems")
.and_then(Value::as_array)
.map(Vec::len)
.unwrap_or(0);
if (0..300).contains(&age) && status == "ok" && problems == 0 {
state.guard_healthy = true;
}
}
}
state
}
fn read_bucket_row(
client: &Client,
server: &str,
host: &str,
bucket: &str,
label: &'static str,
now: DateTime<Utc>,
context: &ContextState,
) -> BucketRow {
let bucket_full = format!("{bucket}_{host}");
let event = latest_event(client, server, &bucket_full).ok().flatten();
let Some(event) = event else {
return BucketRow {
label,
last_id: "0".to_string(),
age: "none".to_string(),
status: classify_bucket_no_events(bucket, context),
};
};
let last_id = event
.get("id")
.map(json_value_to_string)
.unwrap_or_else(|| "0".to_string());
let Some(ts_raw) = event.get("timestamp").and_then(Value::as_str) else {
return BucketRow {
label,
last_id,
age: "?".to_string(),
status: BucketStatus::Unknown,
};
};
let Some(ts) = parse_ts(ts_raw) else {
return BucketRow {
label,
last_id,
age: "?".to_string(),
status: BucketStatus::Unknown,
};
};
let effective_ts = if bucket == "aw-watcher-afk" {
bucket_metadata_end(client, server, &bucket_full).unwrap_or(ts)
} else {
ts
};
let age_sec = (now - effective_ts).num_seconds().max(0);
BucketRow {
label,
last_id,
age: format_age(age_sec),
status: classify_bucket_age(bucket, age_sec, context),
}
}
fn classify_bucket_age(bucket: &str, age_sec: i64, context: &ContextState) -> BucketStatus {
match bucket {
"aw-watcher-window" if context.host_inactive => return BucketStatus::Inactive,
"aw-dlp-endpoint-signals" if context.host_inactive && context.guard_healthy => {
return BucketStatus::Inactive;
}
_ => {}
}
match bucket {
"aw-dlp-incidents" | "aw-dlp-review" | "aw-dlp-rules" | "aw-session-events" => {
if age_sec < 86_400 {
BucketStatus::Fresh
} else {
BucketStatus::EventDriven
}
}
_ if age_sec < 3_600 => BucketStatus::Fresh,
_ if age_sec < 86_400 => BucketStatus::Stale,
_ => BucketStatus::Dead,
}
}
fn classify_bucket_no_events(bucket: &str, context: &ContextState) -> BucketStatus {
match bucket {
"aw-watcher-window" if context.host_inactive => BucketStatus::Inactive,
"aw-dlp-endpoint-signals" if context.host_inactive && context.guard_healthy => {
BucketStatus::Inactive
}
"aw-dlp-incidents" | "aw-dlp-review" | "aw-dlp-rules" | "aw-session-events" => {
BucketStatus::EventDriven
}
_ => BucketStatus::Empty,
}
}
fn summarize(rows: &[BucketRow]) -> Summary {
let mut summary = Summary::default();
for row in rows {
match row.status {
BucketStatus::Fresh | BucketStatus::EventDriven | BucketStatus::Inactive => {
summary.fresh += 1
}
BucketStatus::Stale => summary.stale += 1,
BucketStatus::Dead | BucketStatus::Empty | BucketStatus::Unknown => summary.dead += 1,
}
}
summary
}
fn render_status(colors: &Colors, status: BucketStatus) -> String {
match status {
BucketStatus::Fresh => colors.paint(colors.green, "FRESH"),
BucketStatus::Stale => colors.paint(colors.yellow, "STALE"),
BucketStatus::Dead => colors.paint(colors.red, "DEAD"),
BucketStatus::Empty => colors.paint(colors.red, "EMPTY"),
BucketStatus::EventDriven => colors.paint(colors.cyan, "EVENT-DRIVEN"),
BucketStatus::Inactive => colors.paint(colors.cyan, "INACTIVE"),
BucketStatus::Unknown => colors.paint(colors.red, "?"),
}
}
fn latest_event(client: &Client, server: &str, bucket: &str) -> Result<Option<Value>> {
let url = format!("{server}/api/0/buckets/{bucket}/events?limit=1");
let value = get_json(client, &url)?;
Ok(value.as_array().and_then(|items| items.first()).cloned())
}
fn bucket_metadata_end(client: &Client, server: &str, bucket: &str) -> Option<DateTime<Utc>> {
let url = format!("{server}/api/0/buckets/{bucket}");
let value = get_json(client, &url).ok()?;
value
.pointer("/metadata/end")
.and_then(Value::as_str)
.and_then(parse_ts)
}
fn get_json(client: &Client, url: &str) -> Result<Value> {
client
.get(url)
.header("Connection", "close")
.send()
.with_context(|| format!("GET {url}"))?
.error_for_status()
.with_context(|| format!("GET {url} status"))?
.json::<Value>()
.with_context(|| format!("decode JSON from {url}"))
}
fn check_cors(client: &Client, server: &str) -> u16 {
let url = format!("{server}/api/0/settings/");
client
.get(&url)
.header("Origin", server)
.send()
.map(|response| response.status().as_u16())
.unwrap_or(0)
}
fn tcp_open(host: &str, port: u16, timeout: Duration) -> bool {
let Ok(mut addrs) = (host, port).to_socket_addrs() else {
return false;
};
addrs.any(|addr| TcpStream::connect_timeout(&addr, timeout).is_ok())
}
fn event_timestamp(event: &Value) -> Option<DateTime<Utc>> {
event
.get("timestamp")
.and_then(Value::as_str)
.and_then(parse_ts)
}
fn parse_ts(value: &str) -> Option<DateTime<Utc>> {
DateTime::parse_from_rfc3339(&value.replace('Z', "+00:00"))
.map(|ts| ts.with_timezone(&Utc))
.ok()
}
fn format_age(age_sec: i64) -> String {
if age_sec < 3_600 {
format!("{}m", age_sec / 60)
} else if age_sec < 86_400 {
format!("{}h", age_sec / 3_600)
} else {
format!("{}d", age_sec / 86_400)
}
}
fn json_value_to_string(value: &Value) -> String {
match value {
Value::String(value) => value.clone(),
other => other.to_string(),
}
}
fn env_nonempty(name: &str) -> Option<String> {
std::env::var(name).ok().filter(|value| !value.is_empty())
}
fn parse_env_flag(value: &str) -> bool {
matches!(
value.trim().to_ascii_lowercase().as_str(),
"1" | "true" | "yes" | "on"
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn event_driven_bucket_is_not_dead_when_old() {
assert_eq!(
classify_bucket_age("aw-session-events", 100_000, &ContextState::default()),
BucketStatus::EventDriven
);
}
#[test]
fn inactive_context_suppresses_expected_idle_buckets() {
let context = ContextState {
host_inactive: true,
guard_healthy: true,
};
assert_eq!(
classify_bucket_age("aw-watcher-window", 100_000, &context),
BucketStatus::Inactive
);
assert_eq!(
classify_bucket_no_events("aw-dlp-endpoint-signals", &context),
BucketStatus::Inactive
);
}
#[test]
fn summary_treats_event_and_inactive_as_fresh_class() {
let rows = vec![
BucketRow {
label: "a",
last_id: "0".to_string(),
age: "1m".to_string(),
status: BucketStatus::Fresh,
},
BucketRow {
label: "b",
last_id: "0".to_string(),
age: "none".to_string(),
status: BucketStatus::EventDriven,
},
BucketRow {
label: "c",
last_id: "0".to_string(),
age: "none".to_string(),
status: BucketStatus::Inactive,
},
BucketRow {
label: "d",
last_id: "0".to_string(),
age: "none".to_string(),
status: BucketStatus::Empty,
},
];
let summary = summarize(&rows);
assert_eq!(summary.fresh, 3);
assert_eq!(summary.stale, 0);
assert_eq!(summary.dead, 1);
}
#[test]
fn formats_age_like_legacy_full_check() {
assert_eq!(format_age(59), "0m");
assert_eq!(format_age(3_600), "1h");
assert_eq!(format_age(86_400), "1d");
}
}