# 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 threat intelligence."""

import json
import logging

from secops_mcp.server import get_chronicle_client, server


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

@server.tool()
async def get_threat_intel(
    query: str,
    project_id: str = None,
    customer_id: str = None,
    region: str = None,
) -> str:
    """Get answers to security questions using Chronicle's integrated Gemini model.

    Leverages Gemini to answer general security domain questions or provide specific
    threat intelligence summaries about threat actors, IOCs, CVEs, TTPs, and other
    security topics based on Google's threat intelligence.

    **Workflow Integration:**
    - Use this tool anytime during an investigation to quickly get context or summaries
      on specific threats, vulnerabilities, TTPs, or security concepts encountered
      across any connected security tool (SIEM, SOAR, EDR, TI Platforms, etc.).
    - Can provide valuable background information before diving deeper with more specific
      investigation tools (e.g., entity lookup, event search, file analysis).
    - Useful for understanding the potential impact or nature of a CVE identified by
      any vulnerability management or security posture tool.

    **Use Cases:**
    - "Summarize threat actor APT41."
    - "What is CVE-2024-23897?"
    - "Explain the MITRE ATT&CK technique T1059.001."
    - "What are common persistence mechanisms on Linux?"
    - "Tell me about the LockBit ransomware."

    Args:
        query (str): The security or threat intelligence question to ask Gemini.
        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.

    Returns:
        str: A formatted answer generated by the Gemini model based on the query.
             Returns an error message if the query fails.

    Next Steps (using MCP-enabled tools):
        - Use the provided summary or explanation to inform further investigation steps across relevant tools.
        - If the query was about an indicator (IOC, CVE), consider using entity lookup or SIEM event search tools
          to check for its presence or activity within your environment logs.
        - Correlate the threat intelligence with specific alerts or findings from other security tools (EDR, Network, Cloud).
        - Document relevant findings in the appropriate case management or ticketing system using an MCP tool.
    """
    try:
        logger.info(f'Getting threat intelligence for query: {query}')

        chronicle = get_chronicle_client(project_id, customer_id, region)

        # Call the Gemini method from the SecOps SDK
        response = chronicle.gemini(query)

        # Handle GeminiResponse object
        if hasattr(response, 'get_text_content'):
            # This is a GeminiResponse object, extract text content
            return response.get_text_content()
        elif hasattr(response, 'blocks') and isinstance(response.blocks, list):
            # Handle direct access to blocks if get_text_content isn't available
            text_content = []
            for block in response.blocks:
                if hasattr(block, 'block_type') and hasattr(block, 'content'):
                    if block.block_type == "TEXT":
                        text_content.append(block.content)
            return "\n\n".join(text_content) if text_content else "No text content found in response."
        elif isinstance(response, dict) and 'answer' in response:
            # Legacy format or different API response
            return response.get('answer', 'No answer was provided by the model.')
        elif isinstance(response, str):
            # Direct string response
            return response
        else:
            # If response is in an unexpected format, try to convert it to string
            return json.dumps(response)

    except Exception as e:
        logger.error(f'Error getting threat intelligence: {str(e)}', exc_info=True)
        return f'Error retrieving threat intelligence: {str(e)}'
