analyzer.py

import google.generativeai as genai
import config
import json

SYSTEM_INSTRUCTION = """You are an expert cybersecurity analyst and digital forensics investigator 
specializing in Google Workspace environments. You work as a security operations center (SOC) 
analyst helping IT administrators investigate potential security incidents.

When analyzing logs, you should:
- Identify indicators of compromise (IoC)
- Correlate events across different log types (login, admin, drive, token)
- Assess severity using industry-standard frameworks
- Provide actionable remediation recommendations
- Flag anything that warrants immediate action

Be direct, professional, and thorough. Use bullet points and clear formatting.
When the administrator asks follow-up questions, use the log data you've already been given as context."""

class LogAnalyzer:
    def __init__(self):
        if not config.GEMINI_API_KEY:
            raise ValueError("GEMINI_API_KEY is not set in config.")
        genai.configure(api_key=config.GEMINI_API_KEY)
        self.model = genai.GenerativeModel(
            'gemini-2.5-flash',
            system_instruction=SYSTEM_INSTRUCTION
        )
        self._chat_session = None

    def analyze_logs(self, application_name, logs):
        """
        Analyzes a batch of logs using the Gemini API (one-shot, no conversation).
        Used by the routine scan mode.
        """
        if not logs:
            return "No logs provided for analysis."

        print(f"  Analyzing {len(logs)} {application_name} events with Gemini...")
        prompt = self._construct_scan_prompt(application_name, logs)

        try:
            response = self.model.generate_content(prompt)
            return response.text
        except Exception as e:
            return f"Error during Gemini analysis: {e}"

    def investigate_user(self, user_email, all_logs, user_query):
        """
        Starts a focused investigation on a specific user.
        Opens a chat session so the admin can ask follow-up questions.
        :param user_email: The email being investigated.
        :param all_logs: Dict of {app_name: [events]} for that user.
        :param user_query: The admin's original natural language request.
        :returns: Gemini's initial investigation report (str).
        """
        # Flatten logs into a single summary for the prompt
        log_summary = {}
        total_events = 0
        for app, events in all_logs.items():
            log_summary[app] = events
            total_events += len(events)

        prompt = f"""
An administrator has initiated a security investigation with the following request:

> {user_query}

**Subject:** {user_email}
**Total events collected:** {total_events} across {len(all_logs)} log categories

Below are the Google Workspace audit logs for this user. Perform a thorough forensic analysis.

**Specifically:**
1. **Timeline**: Build a timeline of the user's activity.
2. **Anomalies**: Identify anything suspicious — unusual login locations/times, 
   failed attempts, privilege changes, mass file operations, new OAuth grants, etc.
3. **Risk Assessment**: Provide an overall risk level (Low / Medium / High / Critical) 
   with justification.
4. **Indicators of Compromise**: List any IoCs found.
5. **Recommended Actions**: Provide concrete next steps the admin should take 
   (e.g., reset password, revoke sessions, disable account, review shared files).

If the logs appear normal, state that clearly and explain why.

---
AUDIT LOGS:
{json.dumps(log_summary, indent=2)}
"""
        # Start a chat session so follow-ups have context
        self._chat_session = self.model.start_chat(history=[])
        try:
            response = self._chat_session.send_message(prompt)
            return response.text
        except Exception as e:
            return f"Error during investigation analysis: {e}"

    def chat(self, message):
        """
        Send a follow-up message in the current investigation chat session.
        The model retains full context of the logs and previous conversation.
        """
        if not self._chat_session:
            return "No active investigation session. Start an investigation first."
        try:
            response = self._chat_session.send_message(message)
            return response.text
        except Exception as e:
            return f"Error during follow-up: {e}"

    def end_investigation(self):
        """Ends the current chat session."""
        self._chat_session = None

    def _construct_scan_prompt(self, application_name, logs):
        """Prompt for the routine scan mode (non-conversational)."""
        return f"""
Analyze these recent '{application_name}' audit logs from Google Workspace.
Identify any suspicious activities, anomalies, or potential security risks.

Look for:
- Logins from unexpected or high-risk geographic locations
- Multiple failed login attempts (brute-force)
- Suspicious admin activity (privilege escalation, org unit changes)
- Unusual file sharing, downloading, or deletion (data exfiltration)
- Unusual API usage or token grants

For each finding provide: anomaly type, severity (Low/Medium/High/Critical), 
explanation, and the user or IP involved.

If nothing suspicious is found, state: "No significant security anomalies detected."

Logs:
{json.dumps(logs, indent=2)}
"""