tryaudex_core/

watch.rs

use std::collections::HashSet;

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

use crate::error::{AvError, Result};
use crate::session::{Session, SessionStatus};

/// A single API event observed during a session.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ApiEvent {
    pub timestamp: DateTime<Utc>,
    pub event_name: String,
    pub service: String,
    pub source_ip: Option<String>,
    pub user_agent: Option<String>,
    pub error_code: Option<String>,
    pub error_message: Option<String>,
    pub read_only: bool,
    /// CloudTrail event ID. Used as the dedup key when available; falls back
    /// to a "{timestamp_millis}:{event_name}" composite when absent.
    pub event_id: Option<String>,
}

impl ApiEvent {
    /// Stable dedup key: the CloudTrail event ID when present, otherwise a
    /// composite of millisecond timestamp and event name.
    fn dedup_key(&self) -> String {
        match &self.event_id {
            Some(id) => id.clone(),
            None => format!("{}:{}", self.timestamp.timestamp_millis(), self.event_name),
        }
    }
}

/// Maximum number of events retained in memory to prevent unbounded growth.
const MAX_EVENTS: usize = 10_000;

/// Accumulated watch state for a session.
#[derive(Debug, Default)]
pub struct WatchState {
    pub events: Vec<ApiEvent>,
    /// Fast dedup lookup keyed by `ApiEvent::dedup_key()`.
    seen: HashSet<String>,
    pub last_poll: Option<DateTime<Utc>>,
    pub total_api_calls: usize,
    pub error_count: usize,
    pub services_used: HashSet<String>,
}

impl WatchState {
    /// Merge new events into the state, deduplicating by event ID (falling
    /// back to timestamp+event_name). Caps stored events at `MAX_EVENTS`,
    /// discarding oldest when full.
    pub fn merge(&mut self, new_events: Vec<ApiEvent>) {
        // R6-M12: track the newest CloudTrail event timestamp from this
        // batch and advance `last_poll` to *just past* it, rather than
        // anchoring to local `Utc::now()`. Under clock skew between the
        // local machine and CloudTrail, `Utc::now()` could drift behind
        // the newest event timestamp — each subsequent poll would then
        // re-fetch the same events, dedup them, and observe zero new
        // activity forever. Using `newest_event_ts + 1ms` guarantees
        // monotonic forward progress regardless of clock drift.
        let mut newest_ts: Option<DateTime<Utc>> = None;
        for event in new_events {
            let key = event.dedup_key();
            if self.seen.contains(&key) {
                continue;
            }
            if event.error_code.is_some() {
                self.error_count += 1;
            }
            let svc = event.service.clone();
            self.total_api_calls += 1;
            self.services_used.insert(svc);
            self.seen.insert(key);
            newest_ts = Some(match newest_ts {
                Some(t) if t >= event.timestamp => t,
                _ => event.timestamp,
            });
            self.events.push(event);
        }
        // Evict oldest events if over cap
        if self.events.len() > MAX_EVENTS {
            let excess = self.events.len() - MAX_EVENTS;
            for ev in self.events.drain(..excess) {
                self.seen.remove(&ev.dedup_key());
            }
        }
        let now = Utc::now();
        self.last_poll = Some(match newest_ts {
            Some(newest) => {
                let advanced = newest + chrono::Duration::milliseconds(1);
                if advanced > now {
                    advanced
                } else {
                    now
                }
            }
            None => now,
        });
    }
}

/// Poll CloudTrail LookupEvents for recent activity by an access key.
/// Returns parsed API events since `start_time`.
pub async fn poll_cloudtrail(
    access_key_id: &str,
    start_time: DateTime<Utc>,
    region: Option<&str>,
) -> Result<Vec<ApiEvent>> {
    let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
    if let Some(region) = region {
        loader = loader.region(aws_config::Region::new(region.to_string()));
    }
    let config = loader.load().await;
    let client = aws_sdk_cloudtrail::Client::new(&config);

    let start = aws_sdk_cloudtrail::primitives::DateTime::from_secs(start_time.timestamp());

    let lookup_attr = aws_sdk_cloudtrail::types::LookupAttribute::builder()
        .attribute_key(aws_sdk_cloudtrail::types::LookupAttributeKey::AccessKeyId)
        .attribute_value(access_key_id)
        .build()
        .map_err(|e| AvError::Sts(format!("CloudTrail attribute build error: {}", e)))?;

    // Paginate to collect all events (max 500 to bound memory/API calls).
    let mut all_raw_events = Vec::new();
    let mut next_token: Option<String> = None;
    const MAX_EVENTS: usize = 500;

    loop {
        let mut req = client
            .lookup_events()
            .lookup_attributes(lookup_attr.clone())
            .start_time(start)
            .max_results(50);

        if let Some(ref token) = next_token {
            req = req.next_token(token);
        }

        let result = req
            .send()
            .await
            .map_err(|e| AvError::Sts(format!("CloudTrail LookupEvents error: {}", e)))?;

        all_raw_events.extend(result.events().iter().cloned());

        if all_raw_events.len() >= MAX_EVENTS {
            break;
        }

        match result.next_token() {
            Some(t) if !t.is_empty() => next_token = Some(t.to_string()),
            _ => break,
        }
    }

    let events = all_raw_events
        .iter()
        .filter_map(|e| {
            let event_name = e.event_name()?.to_string();

            // Parse service from event source using the shared CloudTrail mapping
            let service = e
                .event_source()
                .and_then(|s| crate::learn::event_source_to_service(s).map(|svc| svc.to_string()))
                .unwrap_or_else(|| "unknown".to_string());

            let read_only = e
                .read_only()
                .and_then(|r| r.parse::<bool>().ok())
                .unwrap_or(true);

            // Parse error info from CloudTrail event JSON if available
            let (error_code, error_message) = e
                .cloud_trail_event()
                .and_then(|json| serde_json::from_str::<serde_json::Value>(json).ok())
                .map(|v| {
                    let code = v
                        .get("errorCode")
                        .and_then(|c| c.as_str())
                        .map(String::from);
                    let msg = v
                        .get("errorMessage")
                        .and_then(|m| m.as_str())
                        .map(String::from);
                    (code, msg)
                })
                .unwrap_or((None, None));

            let timestamp = e
                .event_time()
                .map(|t| {
                    DateTime::from_timestamp(t.secs(), t.subsec_nanos()).unwrap_or_else(|| {
                        tracing::warn!(
                            secs = t.secs(),
                            "CloudTrail event has unparseable timestamp; falling back to Utc::now()"
                        );
                        Utc::now()
                    })
                })
                .unwrap_or_else(|| {
                    tracing::warn!(
                        "CloudTrail event missing event_time; falling back to Utc::now()"
                    );
                    Utc::now()
                });

            let source_ip = e
                .cloud_trail_event()
                .and_then(|json| serde_json::from_str::<serde_json::Value>(json).ok())
                .and_then(|v| {
                    v.get("sourceIPAddress")
                        .and_then(|s| s.as_str())
                        .map(String::from)
                });

            let user_agent = e
                .cloud_trail_event()
                .and_then(|json| serde_json::from_str::<serde_json::Value>(json).ok())
                .and_then(|v| {
                    v.get("userAgent")
                        .and_then(|s| s.as_str())
                        .map(String::from)
                });

            let event_id = e.event_id().map(String::from);

            Some(ApiEvent {
                timestamp,
                event_name,
                service,
                source_ip,
                user_agent,
                error_code,
                error_message,
                read_only,
                event_id,
            })
        })
        .collect();

    Ok(events)
}

/// Check if a session is still watchable (active and not expired).
pub fn is_watchable(session: &Session) -> bool {
    session.status == SessionStatus::Active && !session.is_expired()
}

/// Format an event as a colored terminal line.
pub fn format_event(event: &ApiEvent) -> String {
    let ts = event.timestamp.format("%H:%M:%S");
    let rw = if event.read_only { "R" } else { "W" };
    let status = if let Some(ref code) = event.error_code {
        format!(" [ERR: {}]", code)
    } else {
        String::new()
    };
    format!(
        "[{}] {} {:<12} {}{}",
        ts, rw, event.service, event.event_name, status
    )
}

/// Format a summary of the current watch state.
pub fn format_summary(state: &WatchState, session: &Session) -> String {
    let mut out = String::new();
    out.push_str(&format!(
        "Session {} | {} API calls | {} errors | {} services | {}s remaining\n",
        session.short_id(),
        state.total_api_calls,
        state.error_count,
        state.services_used.len(),
        session.remaining_seconds()
    ));
    if !state.services_used.is_empty() {
        let mut svcs: Vec<&str> = state.services_used.iter().map(String::as_str).collect();
        svcs.sort_unstable();
        out.push_str(&format!("Services: {}\n", svcs.join(", ")));
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_event(name: &str, service: &str, error: Option<&str>) -> ApiEvent {
        ApiEvent {
            timestamp: Utc::now(),
            event_name: name.to_string(),
            service: service.to_string(),
            source_ip: Some("10.0.0.1".to_string()),
            user_agent: Some("aws-cli/2.0".to_string()),
            error_code: error.map(String::from),
            error_message: None,
            read_only: true,
            event_id: None,
        }
    }

    #[test]
    fn test_watch_state_merge() {
        let mut state = WatchState::default();
        let events = vec![
            make_event("GetObject", "s3", None),
            make_event("ListBuckets", "s3", None),
        ];
        state.merge(events);
        assert_eq!(state.total_api_calls, 2);
        assert_eq!(state.error_count, 0);
        assert!(state.services_used.contains("s3"));
        assert_eq!(state.services_used.len(), 1);
    }

    #[test]
    fn test_watch_state_dedup() {
        let mut state = WatchState::default();
        let ts = Utc::now();
        let event = ApiEvent {
            timestamp: ts,
            event_name: "GetObject".to_string(),
            service: "s3".to_string(),
            source_ip: None,
            user_agent: None,
            error_code: None,
            error_message: None,
            read_only: true,
            event_id: None,
        };
        state.merge(vec![event.clone()]);
        state.merge(vec![event]);
        assert_eq!(state.total_api_calls, 1); // deduped
    }

    #[test]
    fn test_watch_state_error_count() {
        let mut state = WatchState::default();
        let events = vec![
            make_event("GetObject", "s3", None),
            make_event("PutObject", "s3", Some("AccessDenied")),
        ];
        state.merge(events);
        assert_eq!(state.error_count, 1);
    }

    #[test]
    fn test_watch_state_multi_service() {
        let mut state = WatchState::default();
        let events = vec![
            make_event("GetObject", "s3", None),
            make_event("GetFunction", "lambda", None),
            make_event("ListBuckets", "s3", None),
        ];
        state.merge(events);
        assert_eq!(state.services_used.len(), 2);
        assert!(state.services_used.contains("s3"));
        assert!(state.services_used.contains("lambda"));
    }

    #[test]
    fn test_format_event() {
        let event = make_event("GetObject", "s3", None);
        let line = format_event(&event);
        assert!(line.contains("s3"));
        assert!(line.contains("GetObject"));
        assert!(line.contains("R")); // read-only
    }

    #[test]
    fn test_format_event_with_error() {
        let event = make_event("PutObject", "s3", Some("AccessDenied"));
        let line = format_event(&event);
        assert!(line.contains("AccessDenied"));
    }

    #[test]
    fn test_format_summary() {
        let mut state = WatchState::default();
        state.merge(vec![
            make_event("GetObject", "s3", None),
            make_event("GetFunction", "lambda", None),
        ]);
        let session = crate::session::Session::new(
            std::time::Duration::from_secs(900),
            None,
            crate::policy::ScopedPolicy::from_allow_str("s3:GetObject").unwrap(),
            "arn:aws:iam::123456789012:role/Test".to_string(),
            vec!["test".to_string()],
        );
        let summary = format_summary(&state, &session);
        assert!(summary.contains("2 API calls"));
        assert!(summary.contains("2 services"));
    }
}