#!/usr/bin/env python
#
# tyr/tools/get_tyr_user_risk.py - MCP Tool for calculating user risk scores
#
import os
from google.cloud import bigquery
from datetime import datetime, timedelta

RISK_WEIGHTS = {
    "google.cloud.secretmanager.v1.SecretManagerService.AccessSecretVersion": 10,
    "google.iam.admin.v1.SetIamPolicy": 20,
    "google.storage.v1.Storage.DeleteBucket": 25,
    "google.login.LoginService.LoginFailure": 5,
}

def get_tyr_user_risk(p: dict) -> dict:
    """Calculates a risk score for a user based on recent audit log activity."""
    project_id = os.environ.get("GOOGLE_CLOUD_PROJECT", "propane-will-491900-m5")
    principal_email = p.get("principal_email")
    caller_ip = p.get("caller_ip")
    hours = p.get("hours", 24)

    if not (principal_email or caller_ip):
        raise ValueError("Missing required parameter: 'principal_email' or 'caller_ip'")

    client = bigquery.Client()
    table_id = f"{project_id}.tyr_audit_logs.cloudaudit_googleapis_com_activity"
    
    where_clauses = [f"timestamp >= TIMESTAMP_SUB(CURRENT_TIMESTAMP(), INTERVAL {hours} HOUR)"]
    if principal_email:
        where_clauses.append(f"protopayload_auditlog.authenticationInfo.principalEmail = '{principal_email}'")
    if caller_ip:
        where_clauses.append(f"protopayload_auditlog.requestMetadata.callerIp = '{caller_ip}'")
    
    query = f"""
        SELECT
            timestamp,
            protopayload_auditlog.methodName,
            protopayload_auditlog.status.message as status_message
        FROM `{table_id}`
        WHERE {" AND ".join(where_clauses)}
    """

    try:
        query_job = client.query(query)
        results = query_job.result()
        
        risk_score = 0
        findings = []
        
        for row in results:
            weight = RISK_WEIGHTS.get(row.methodName, 0) # Default to 0 for non-risky methods
            if row.status_message and "authentication failed" in row.status_message.lower():
                weight = RISK_WEIGHTS.get("google.login.LoginService.LoginFailure", 5)
            
            if weight > 0:
                risk_score += weight
                findings.append({
                    "timestamp": row.timestamp.isoformat(),
                    "method": row.methodName,
                    "status": row.status_message,
                    "weight": weight
                })

        return {
            "status": "success",
            "principal_email": principal_email,
            "caller_ip": caller_ip,
            "time_window_hours": hours,
            "calculated_risk_score": risk_score,
            "findings_count": len(findings),
            "findings": findings
        }

    except Exception as e:
        print(f"An error occurred while querying audit logs: {e}")
        raise