agent.py

"""
Agentic Security Investigator — powered by Gemini Function Calling.

Instead of the admin deciding which logs to pull, Gemini autonomously:
  1. Interprets the admin's natural-language request.
  2. Calls the right workspace functions (fetch user logs, fetch org logs, etc.).
  3. Analyzes the returned data.
  4. Decides if more data is needed and fetches it.
  5. Delivers a final investigation report.

The admin just describes the situation and the AI does the rest.
"""

import json
import google.generativeai as genai
from google.protobuf.struct_pb2 import Struct
from workspace_client import WorkspaceClient
import config

# ── Colors ──────────────────────────────────────────────────────
CYAN   = "\033[0;36m"
GREEN  = "\033[0;32m"
YELLOW = "\033[1;33m"
RED    = "\033[0;31m"
BOLD   = "\033[1m"
DIM    = "\033[2m"
NC     = "\033[0m"

# ── Tool definitions (what Gemini can call) ─────────────────────
# Using the manual function-declaration dict format for google.generativeai

TOOL_DECLARATIONS = [
    {
        "name": "fetch_user_logs",
        "description": (
            "Fetch Google Workspace audit logs for a SPECIFIC user. "
            "Returns login, admin, drive, and token events for that user. "
            "Use this when investigating a particular user."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "user_email": {
                    "type": "string",
                    "description": "The email address of the user to investigate."
                },
                "days_back": {
                    "type": "integer",
                    "description": "How many days of history to fetch. Default 7. Use 30 for deeper investigations."
                },
                "application": {
                    "type": "string",
                    "enum": ["login", "admin", "drive", "token", "user_accounts",
                             "gmail", "calendar", "chat", "meet", "mobile",
                             "groups_enterprise", "rules", "saml", "all"],
                    "description": (
                        "Which log type to fetch. Options: "
                        "'login' (sign-ins), 'admin' (admin changes), 'drive' (file activity), "
                        "'token' (OAuth grants), 'user_accounts' (email settings/forwarding/delegates), "
                        "'gmail' (email send/receive/spam), 'calendar', 'chat', 'meet', "
                        "'mobile' (device mgmt), 'groups_enterprise', 'rules' (DLP), "
                        "'saml' (SSO), or 'all' to fetch everything."
                    )
                }
            },
            "required": ["user_email"]
        }
    },
    {
        "name": "fetch_org_logs",
        "description": (
            "Fetch Google Workspace audit logs for the ENTIRE organization. "
            "Returns events for all users. Use this for org-wide sweeps or "
            "when no specific user is mentioned."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "days_back": {
                    "type": "integer",
                    "description": "How many days of history to fetch. Default 1."
                },
                "application": {
                    "type": "string",
                    "enum": ["login", "admin", "drive", "token", "user_accounts",
                             "gmail", "calendar", "chat", "meet", "mobile",
                             "groups_enterprise", "rules", "saml"],
                    "description": (
                        "Which audit log category to fetch. Use 'gmail' or 'user_accounts' "
                        "for email-related investigations."
                    )
                }
            },
            "required": ["application"]
        }
    },
    {
        "name": "list_available_log_types",
        "description": (
            "Returns the list of audit log categories available to query. "
            "Call this if you need to know what data sources are available."
        ),
        "parameters": {
            "type": "object",
            "properties": {}
        }
    }
]

SYSTEM_PROMPT = """You are an autonomous cybersecurity AI agent embedded in a Google Workspace 
security operations center. You have direct access to the organization's audit logs through 
function calls.

When an administrator describes a security concern, you should:
1. THINK about what data you need to investigate the issue.
2. CALL the appropriate functions to fetch the relevant logs.
3. ANALYZE the returned data thoroughly.
4. If you need MORE data to complete the investigation, call more functions.
5. Provide a comprehensive security report.

You can call multiple functions in sequence. For example, if someone reports a compromised 
account, you might:
- First fetch their login logs to check for suspicious sign-ins.
- Then fetch their drive logs to check for data exfiltration.
- Then fetch their token logs to check for unauthorized app access.
- Then fetch org-wide admin logs to see if privileges were escalated.

ALWAYS fetch data before making conclusions. Never guess — use the tools.

When delivering findings, use this structure:
- **Summary**: One-paragraph overview.
- **Timeline**: Key events in chronological order.
- **Risk Level**: Low / Medium / High / Critical with justification.
- **Indicators of Compromise**: Specific suspicious findings.
- **Recommended Actions**: Concrete steps the admin should take.

Be thorough but concise. The administrator is technical and expects professional SOC-level analysis."""


class AgentInvestigator:
    """
    An agentic investigation engine that lets Gemini autonomously decide 
    which Workspace logs to fetch and analyze.
    """

    def __init__(self):
        if not config.GEMINI_API_KEY:
            raise ValueError("GEMINI_API_KEY is not set.")
        genai.configure(api_key=config.GEMINI_API_KEY)

        self.workspace = WorkspaceClient()
        self._model_index = 0
        self._history = []         # preserved across model switches
        self._build_model()

    def _build_model(self):
        """Build (or rebuild) the GenerativeModel with the current model index."""
        model_name = config.MODELS[self._model_index]
        print(f"  {DIM}Using model: {model_name}{NC}")
        self.model = genai.GenerativeModel(
            model_name,
            tools=[{"function_declarations": TOOL_DECLARATIONS}],
            system_instruction=SYSTEM_PROMPT
        )
        self.chat = self.model.start_chat(
            history=list(self._history),
            enable_automatic_function_calling=False
        )

    def start_session(self):
        """Start a new agentic chat session (resets history)."""
        self._model_index = 0
        self._history = []
        self._build_model()
        self._turn_count = 0

    def _switch_model(self):
        """Try the next model in the fallback list. Returns True if switched."""
        self._model_index += 1
        if self._model_index >= len(config.MODELS):
            return False
        print(f"  {YELLOW}↻ Switching model...{NC}")
        self._build_model()
        return True

    def _send_with_retry(self, message, max_retries=3):
        """Send a message with retry logic + automatic model fallback."""
        import time
        while True:
            for attempt in range(max_retries):
                try:
                    response = self.chat.send_message(message)
                    # Save history on success (for model-switch resilience)
                    self._history = list(self.chat.history)
                    return response
                except Exception as e:
                    error_str = str(e)
                    if "429" in error_str or "quota" in error_str.lower():
                        wait = 2 ** attempt * 2  # 2s, 4s, 8s
                        print(f"  {YELLOW}⏳ Rate limited — waiting {wait}s before retry ({attempt+1}/{max_retries})...{NC}")
                        time.sleep(wait)
                    else:
                        raise e

            # All retries exhausted — try next model
            if not self._switch_model():
                raise Exception(
                    "All models rate-limited. Wait a minute and try again."
                )

    def process_query(self, user_query):
        """
        Send a query to the agent and run the agentic loop.
        The agent will autonomously call functions until it has enough data 
        to deliver a final answer.
        """
        if not self.chat:
            self.start_session()

        self._turn_count += 1
        print(f"\n{DIM}Agent is thinking...{NC}")

        try:
            response = self._send_with_retry(user_query)
        except Exception as e:
            return f"Error communicating with Gemini: {e}"

        # Agentic loop: keep executing function calls until the model returns text
        max_iterations = 10  # safety limit
        iteration = 0

        while iteration < max_iterations:
            iteration += 1

            # Check if the response contains function calls
            function_calls = self._extract_function_calls(response)

            if not function_calls:
                # No more function calls — the model has a final answer
                break

            # Execute each function call and collect results
            function_responses = []
            for fc in function_calls:
                fname = fc.name
                fargs = dict(fc.args) if fc.args else {}

                print(f"  {GREEN}⚡ Agent calling:{NC} {BOLD}{fname}{NC}({self._format_args(fargs)})")

                result = self._execute_function(fname, fargs)

                # Build the function response part
                function_responses.append(
                    genai.protos.Part(
                        function_response=genai.protos.FunctionResponse(
                            name=fname,
                            response=self._to_proto_struct({"result": result})
                        )
                    )
                )

            # Send function results back to the model
            print(f"  {DIM}Agent analyzing results...{NC}")
            try:
                response = self._send_with_retry(
                    genai.protos.Content(parts=function_responses)
                )
            except Exception as e:
                return f"Error in agentic loop: {e}"

        # Extract the final text response
        if response.candidates and response.candidates[0].content.parts:
            final_parts = []
            for part in response.candidates[0].content.parts:
                if hasattr(part, 'text') and part.text:
                    final_parts.append(part.text)
            return "\n".join(final_parts) if final_parts else "Agent completed but produced no output."
        
        return "Agent completed but produced no output."

    def _extract_function_calls(self, response):
        """Extract function call parts from a model response."""
        calls = []
        if response.candidates and response.candidates[0].content.parts:
            for part in response.candidates[0].content.parts:
                if hasattr(part, 'function_call') and part.function_call.name:
                    calls.append(part.function_call)
        return calls

    def _execute_function(self, name, args):
        """Execute a workspace function by name and return the result."""
        # Gemini sometimes passes days_back as a float (e.g., 30.0)
        if 'days_back' in args:
            args['days_back'] = int(args['days_back'])

        if name == "fetch_user_logs":
            return self._fn_fetch_user_logs(**args)
        elif name == "fetch_org_logs":
            return self._fn_fetch_org_logs(**args)
        elif name == "list_available_log_types":
            return self._fn_list_log_types()
        else:
            return {"error": f"Unknown function: {name}"}

    # ── Log summarization (reduces token count) ─────────────────

    def _summarize_events(self, events, max_events=50):
        """
        Strip bulky nested fields from log events to reduce token usage.
        Keeps only the security-relevant fields.
        """
        summarized = []
        for event in events[:max_events]:
            summary = {}
            # Keep key fields
            for key in ('id', 'actor', 'ipAddress', 'events', 'kind'):
                if key in event:
                    val = event[key]
                    # For 'events', flatten to just name + parameters
                    if key == 'events' and isinstance(val, list):
                        summary['events'] = []
                        for ev in val:
                            flat = {"name": ev.get("name", "")}
                            params = ev.get("parameters", [])
                            if params:
                                flat["parameters"] = {
                                    p.get("name", ""): (
                                        p.get("value") or p.get("boolValue") or
                                        p.get("intValue") or p.get("multiValue", "")
                                    )
                                    for p in params[:10]  # cap params per event
                                }
                            summary['events'].append(flat)
                    else:
                        summary[key] = val
            summarized.append(summary)

        result = summarized
        if len(events) > max_events:
            result.append({"_note": f"Showing {max_events} of {len(events)} total events"})
        return result

    # ── Function implementations ────────────────────────────────

    def _fn_fetch_user_logs(self, user_email, days_back=7, application="all"):
        """Fetch logs for a specific user."""
        if application == "all":
            logs = {}
            total = 0
            for app in config.APPLICATIONS_TO_MONITOR:
                events = self.workspace.fetch_logs(
                    app, user_key=user_email, days_back=days_back
                )
                logs[app] = self._summarize_events(events)
                total += len(events)
            return {
                "user": user_email,
                "days_back": days_back,
                "total_events": total,
                "logs": logs
            }
        else:
            events = self.workspace.fetch_logs(
                application, user_key=user_email, days_back=days_back
            )
            return {
                "user": user_email,
                "application": application,
                "days_back": days_back,
                "total_events": len(events),
                "logs": self._summarize_events(events)
            }

    def _fn_fetch_org_logs(self, application, days_back=1):
        """Fetch org-wide logs for a specific application."""
        events = self.workspace.fetch_logs(
            application, user_key='all', days_back=days_back
        )
        return {
            "application": application,
            "days_back": days_back,
            "total_events": len(events),
            "logs": self._summarize_events(events)
        }

    def _fn_list_log_types(self):
        """List available log types."""
        return {
            "available_log_types": config.APPLICATIONS_TO_MONITOR,
            "descriptions": {
                "login": "User sign-in events, failed attempts, suspicious locations",
                "admin": "Admin console changes, privilege modifications, org settings",
                "drive": "File sharing, downloads, external sharing, deletions",
                "token": "OAuth token grants, third-party app authorizations",
                "user_accounts": "Email forwarding rules, delegates, password changes, POP/IMAP",
                "gmail": "Email send/receive activity, spam, phishing reports",
                "calendar": "Calendar sharing, event modifications",
                "chat": "Google Chat messages and spaces",
                "meet": "Meet calls, participants, recordings",
                "mobile": "Mobile device activity, wipes, policy changes",
                "groups_enterprise": "Group membership changes, settings",
                "rules": "DLP rule triggers and actions",
                "saml": "SAML SSO login events"
            }
        }

    # ── Helpers ──────────────────────────────────────────────────

    def _format_args(self, args):
        """Format function args for display."""
        parts = []
        for k, v in args.items():
            parts.append(f"{k}={repr(v)}")
        return ", ".join(parts)

    def _to_proto_struct(self, data):
        """Convert a Python dict to a protobuf Struct for function responses."""
        # We need to serialize to JSON-safe format first
        # (the API expects simple types, not complex nested objects with large log data)
        safe_data = json.loads(json.dumps(data, default=str))
        struct = Struct()
        struct.update(safe_data)
        return struct