OSVauco/tyr/tools/get_tyr_forecast.py

81 lines
2.8 KiB
Python

#!/usr/bin/env python
#
# tyr/tools/get_tyr_forecast.py - MCP Tool for anomaly detection
#
import os
from google.cloud import bigquery
import statistics
def get_tyr_forecast(p: dict) -> dict:
"""Analyzes audit log volume to find statistical anomalies."""
project_id = os.environ.get(
"GOOGLE_CLOUD_PROJECT", "propane-will-491900-m5")
hours_to_check = p.get("hours_to_check", 24)
std_dev_threshold = p.get("std_dev_threshold", 3.0)
client = bigquery.Client()
dataset_id = "tyr_audit_logs"
# This assumes a table partitioned by day, which is standard for log sinks.
# It queries the last 7 days to build a baseline.
query = f"""
SELECT
TIMESTAMP_TRUNC(timestamp, HOUR) as hour,
COUNT(*) as event_count
FROM `{project_id}.{dataset_id}.cloudaudit_googleapis_com_activity_*`
WHERE _TABLE_SUFFIX BETWEEN
FORMAT_DATE('%Y%m%d', DATE_SUB(CURRENT_DATE(), INTERVAL 8 DAY)) AND
FORMAT_DATE('%Y%m%d', DATE_SUB(CURRENT_DATE(), INTERVAL 1 DAY))
GROUP BY 1
ORDER BY 1
"""
try:
query_job = client.query(query)
results = query_job.result()
hourly_counts = [row.event_count for row in results]
if len(hourly_counts) < 2:
return {"status": "error", "message": "Not enough historical data to calculate a baseline."}
# Calculate baseline mean and standard deviation
mean = statistics.mean(hourly_counts)
stdev = statistics.stdev(hourly_counts)
# Now, check the most recent N hours
query_recent = f"""
SELECT
TIMESTAMP_TRUNC(timestamp, HOUR) as hour,
COUNT(*) as event_count
FROM `{project_id}.{dataset_id}.cloudaudit_googleapis_com_activity_*`
WHERE timestamp >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL {hours_to_check} HOUR)
GROUP BY 1
ORDER BY 1 DESC
"""
query_job_recent = client.query(query_recent)
recent_results = query_job_recent.result()
anomalies = []
for row in recent_results:
if abs(row.event_count - mean) > (stdev * std_dev_threshold):
anomaly = {
"timestamp": row.hour.isoformat(),
"event_count": row.event_count,
"baseline_mean": round(mean, 2),
"baseline_stdev": round(stdev, 2),
"deviation": round((row.event_count - mean) / stdev, 2)
}
anomalies.append(anomaly)
if anomalies:
return {"status": "anomalies_found", "anomalies": anomalies}
else:
return {"status": "success", "message": f"No anomalies found in the last {hours_to_check} hours."}
except Exception as e:
print(f"An error occurred: {e}")
raise