# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Security Operations MCP tools for security alerts."""

import json
import logging
from datetime import datetime, timedelta, timezone

from typing import Any, Dict, Optional, Literal, Union
from secops_mcp.server import get_chronicle_client, server


# Configure logging
logger = logging.getLogger('secops-mcp')

@server.tool()
async def get_security_alerts(
    project_id: str = None,
    customer_id: str = None,
    hours_back: int = 24,
    max_alerts: int = 10,
    status_filter: str = 'feedback_summary.status != "CLOSED"',
    region: str = None,
) -> str:
    """Get security alerts directly from Chronicle SIEM.

    Retrieves a list of recent security alerts generated within Chronicle, based on
    detection rules or other alert sources configured in the SIEM.

    **Workflow Integration:**
    - Use this for direct monitoring of SIEM alert activity, potentially identifying
      issues before they are ingested or processed by other platforms (e.g., SOAR).
    - Can be used as an initial step to get a sense of recent high-priority events
      directly from the source SIEM.
    - Contrast this with tools that list alerts associated with a specific case
      in a case management or SOAR system.

    **Use Cases:**
    - Get a quick overview of recent, non-closed alerts generated by the SIEM.
    - Monitor for specific high-severity alerts or rule triggers.
    - Check for SIEM alerts that might not have corresponding cases yet in other systems.

    Args:
        project_id (Optional[str]): Google Cloud project ID. Defaults to environment configuration.
        customer_id (Optional[str]): Chronicle customer ID. Defaults to environment configuration.
        hours_back (int): How many hours to look back for alerts. Defaults to 24.
        max_alerts (int): Maximum number of alerts to return. Defaults to 10.
        status_filter (str): Query string to filter alerts by status (e.g., 'feedback_summary.status != "CLOSED"').
                             Defaults to excluding closed alerts.
        region (Optional[str]): Chronicle region (e.g., "us", "europe"). Defaults to environment configuration.

    Returns:
        str: A formatted string summarizing the retrieved security alerts, including rule name,
             creation time, status, severity, and associated case ID (if available).
             Returns 'No security alerts found...' if none match the criteria.

    Next Steps (using MCP-enabled tools):
        - Analyze the returned alerts for priority and relevance.
        - For high-priority alerts, check if a corresponding case exists in your case management/SOAR system.
        - If no Alerts are found, expand the filter for this tool by increasing the max_alerts incremently until you are confident there are no recent Alerts
        - If no case exists, consider creating one or initiating investigation directly.
        - Use entity lookup tools (like `lookup_entity`) on indicators found within the alert details.
        - Use SIEM event search tools (like `search_security_events`) to find related raw logs.
        - Correlate alert information with findings from other security tools (EDR, Cloud Posture, TI) via their MCP tools.
    """
    try:
        chronicle = get_chronicle_client(project_id, customer_id, region)

        end_time = datetime.now(timezone.utc)
        start_time = end_time - timedelta(hours=hours_back)

        alert_response = chronicle.get_alerts(
            start_time=start_time,
            end_time=end_time,
            snapshot_query=status_filter,
            max_alerts=max_alerts,
        )

        # The response format depends on the secops library version
        # Try to handle both formats
        if isinstance(alert_response, dict):
            alert_list = alert_response.get('alerts', {}).get('alerts', [])
        else:
            # Might be a direct list of alerts in the standard library
            alert_list = alert_response if isinstance(alert_response, list) else []

        if not alert_list:
            return 'No security alerts found for the specified time range.'

        result = f'Found {len(alert_list)} security alerts:\n\n'

        for i, alert in enumerate(alert_list, 1):
            # Try to access fields with different possible structures
            rule_name = None
            if (
                'detection' in alert
                and isinstance(alert['detection'], list)
                and len(alert['detection']) > 0
            ):
                rule_name = alert['detection'][0].get('ruleName', 'Unknown Rule')
            else:
                rule_name = alert.get('ruleName', 'Unknown Rule')

            created_time = alert.get('createdTime', 'Unknown')

            # Try different possible status field paths
            status = 'Unknown'
            if 'feedbackSummary' in alert and isinstance(
                alert['feedbackSummary'], dict
            ):
                status = alert['feedbackSummary'].get('status', 'Unknown')
            elif 'status' in alert:
                status = alert.get('status', 'Unknown')

            # Try different possible severity field paths
            severity = 'Unknown'
            if 'feedbackSummary' in alert and isinstance(
                alert['feedbackSummary'], dict
            ):
                severity = alert['feedbackSummary'].get('severityDisplay', 'Unknown')
            elif 'severity' in alert:
                severity = alert.get('severity', 'Unknown')

            result += f'Alert {i}:\n'
            result += f'Rule: {rule_name}\n'
            result += f'Created: {created_time}\n'
            result += f'Status: {status}\n'
            result += f'Severity: {severity}\n'

            # Add case information if available
            case_name = alert.get('caseName')
            if case_name:
                result += f'Associated Case: {case_name}\n'

            result += '\n'

        return json.dumps(result)
    except Exception as e:
        return f'Error retrieving security alerts: {str(e)}'

@server.tool()
async def get_security_alert_by_id(
    project_id: str = None,
    customer_id: str = None,
    region: str = None,
    alert_id: str = None,
    include_detections: bool = True
) -> str:
    """Get security alert by ID directly from Chronicle SIEM.

    Gets an alert by ID

    **Workflow Integration:**
    - Use this for direct monitoring of SIEM alert activity, potentially identifying
      issues before they are ingested or processed by other platforms (e.g., SOAR).
    - Can be used as an initial step to get a sense of recent high-priority events
      directly from the source SIEM.
    - Contrast this with tools that list alerts associated with a specific case
      in a case management or SOAR system.

    **Use Cases:**
    - View a specific Alert
    - Monitor for specific high-severity alerts or rule triggers.
    - Check for SIEM alerts that might not have corresponding cases yet in other systems.
    - May need to get this so you know which Alert to update

    Args:
        project_id (Optional[str]): Google Cloud project ID. Defaults to environment configuration.
        customer_id (Optional[str]): Chronicle customer ID. Defaults to environment configuration.
        region (Optional[str]): Chronicle region (e.g., "us", "europe"). Defaults to environment configuration.
        alert_id (Optional[str]): The unique identifier of the alert to retrieve.
        include_detections (bool): Whether to include detection details in the response. Defaults to True.

    Returns:
        str: A formatted string summarizing the retrieved security alerts, including rule name,
             creation time, status, severity, and associated case ID (if available).
             Returns 'No security alerts found...' if none match the criteria.

    Next Steps (using MCP-enabled tools):
        - Analyze the returned alerts for priority and relevance.
        - For high-priority alerts, check if a corresponding case exists in your case management/SOAR system.
        - If no case exists, consider creating one or initiating investigation directly.
        - Use entity lookup tools (like `lookup_entity`) on indicators found within the alert details.
        - Use SIEM event search tools (like `search_security_events`) to find related raw logs.
        - Correlate alert information with findings from other security tools (EDR, Cloud Posture, TI) via their MCP tools.
    """

    try:
        chronicle = get_chronicle_client(project_id, customer_id, region)
        response = chronicle.get_alert(alert_id, include_detections)
    except Exception as e:
        return f'Error retrieving security alert for {alert_id}: {str(e)}'

    return json.dumps(response)

@server.tool()
async def do_update_security_alert(
    project_id: str = None,
    customer_id: str = None,
    region: str = None,
    alert_id: str = None,
    reason: Optional[str] = None,
    priority: Optional[str] = None,
    status: Optional[str] = None,
    verdict: Optional[str] = None,
    severity: Optional[int] = None,
    comment: Optional[Union[str, Literal[""]]] = None,
    root_cause: Optional[Union[str, Literal[""]]] = None
) -> str:
    """
        Update security alert attributes directly in Chronicle SIEM.

Modifies specific fields of an existing security alert within Chronicle based on its ID. This function allows for updates to an alert's status, severity, verdict, assigned scores, comments, and other metadata. This is typically performed after an investigation, triage, or automated analysis provides new insights or conclusions about the alert. At least one of the optional fields related to alert attributes (e.g., status, severity, comment) should be provided to perform a meaningful update.

**Workflow Integration:**
-   Utilize when SOAR is not a core technology the investigator uses
-   Utilize after an initial investigation of an alert (e.g., using `get_chronicle_alert_details` or other analytical tools) to record findings or change its state.
-   Incorporate into automated or semi-automated triage workflows where alert properties are updated based on enrichment data, external threat intelligence, or predefined logic.
-   Use to reflect the outcome of a manual investigation, such as marking an alert as a false positive, confirming it as a true positive, adjusting its severity, or setting a final verdict.
-   Can be used to synchronize alert states between Chronicle SIEM and external systems like SOAR platforms or case management tools, if this tool is part of such an integration.
-   Helps in maintaining an accurate and up-to-date view of alert lifecycle and analyst findings directly within Chronicle.

**Use Cases:**
-   SOAR is not used for alert disposition
-   Change the lifecycle status of an alert (e.g., from "NEW" to "REVIEWED", or to "CLOSED"). Valid statuses include: "STATUS_UNSPECIFIED", "NEW", "REVIEWED", "CLOSED", "OPEN".
-   Adjust the severity (as an integer value, 0-100) or priority (e.g., "PRIORITY_LOW", "PRIORITY_MEDIUM", "PRIORITY_HIGH") of an alert based on new information or impact assessment. Valid priorities include: "PRIORITY_UNSPECIFIED", "PRIORITY_INFO", "PRIORITY_LOW", "PRIORITY_MEDIUM", "PRIORITY_HIGH", "PRIORITY_CRITICAL".
-   Set a definitive verdict on an alert (e.g., "TRUE_POSITIVE", "FALSE_POSITIVE") after detailed analysis. Valid verdicts include: "VERDICT_UNSPECIFIED", "TRUE_POSITIVE", "FALSE_POSITIVE".
-   Add or update investigative comments, root cause analysis details, or reasons for the alert's disposition. Providing an empty string can clear existing comments or root cause.
-   Update an alert's confidence score or risk score based on corroborating evidence or lack thereof.

Args:
    alert_id (str): The unique ID of the Chronicle security alert to update. This is a required identifier.
    project_id (Optional[str]): Google Cloud project ID associated with the Chronicle instance. Defaults to environment configuration if not provided.
    customer_id (Optional[str]): The Chronicle customer ID. Defaults to environment configuration if not provided.
    region (Optional[str]): The Google Cloud region where the Chronicle instance is hosted (e.g., "us", "europe"). Defaults to environment configuration if not provided.
    reason: Reason for closing an alert. Valid values:
        - "REASON_UNSPECIFIED"
        - "REASON_NOT_MALICIOUS"
        - "REASON_MALICIOUS"
        - "REASON_MAINTENANCE"
    priority: Alert priority. Valid values:
        - "PRIORITY_UNSPECIFIED"
        - "PRIORITY_INFO"
        - "PRIORITY_LOW"
        - "PRIORITY_MEDIUM"
        - "PRIORITY_HIGH"
        - "PRIORITY_CRITICAL"
    status: Alert status. Valid values:
        - "STATUS_UNSPECIFIED"
        - "NEW"
        - "REVIEWED"
        - "CLOSED"
        - "OPEN"
    verdict: Verdict on the alert. Valid values:
        - "VERDICT_UNSPECIFIED"
        - "TRUE_POSITIVE"
        - "FALSE_POSITIVE"
    risk_score: Risk score [0-100] of the alert
    severity: Severity score [0-100] of the alert
    comment: Analyst comment (empty string is valid to clear)
    root_cause: Alert root cause (empty string is valid to clear)

Returns:
    str: A confirmation message indicating whether the alert was updated successfully, potentially including a summary of the changes or the updated alert ID. Returns an error message if the update fails.

Next Steps (using MCP-enabled tools):
    - After updating, retrieve the alert again using `get_chronicle_alert_details` with the `alert_id` to verify that the changes have been applied correctly.
    - If the alert status was changed to indicate resolution or a false positive, consider if any adjustments to detection rules or alert thresholds are needed to improve accuracy or reduce noise. Is there an opportunity to create a rule exclusion?
    - Ensure that any corresponding ticket or case in an external case management or SOAR system is updated to reflect the changes made in Chronicle.
    - Communicate significant updates (e.g., confirmed breach, critical false positive) to relevant teams or stakeholders as per incident response procedures.
    """
    try:
        chronicle = get_chronicle_client(project_id, customer_id, region)
        response = chronicle.update_alert(alert_id, reason=reason, status=status, verdict=verdict, comment=comment, root_cause=root_cause, priority=priority, severity=severity)
    except Exception as e:
        return f'Error retrieving security alert for {alert_id}: {str(e)}'

    return json.dumps(response)
