diff --git a/run-with-google-adk/console-mcp-agent.py b/run-with-google-adk/console-mcp-agent.py index febbd8ba..03213ccc 100644 --- a/run-with-google-adk/console-mcp-agent.py +++ b/run-with-google-adk/console-mcp-agent.py @@ -86,6 +86,8 @@ async def get_all_tools(): secops_tools = [] gti_tools = [] secops_soar_tools = [] + scc_tools = [] + compliance_tools = [] exit_stack = AsyncExitStack() uv_dir_prefix="../server" @@ -153,9 +155,22 @@ async def get_all_tools(): ),async_exit_stack=exit_stack ) + if os.environ.get("LOAD_COMPLIANCE_MCP") == "Y": + compliance_tools, exit_stack = await MCPToolset.from_server( + connection_params=StdioServerParameters( + command='uv', + args=[ "--directory", + uv_dir_prefix + "/compliance/", + "run", + "compliance_mcp.py" + ], + ),async_exit_stack=exit_stack + ) + logging.info("MCP Toolsets created successfully.") - return secops_tools+gti_tools+secops_soar_tools+scc_tools, exit_stack + all_tools = secops_tools + gti_tools + secops_soar_tools + scc_tools + compliance_tools + return all_tools, exit_stack async def get_agent_async(): """Gets tools from MCP Server.""" diff --git a/run-with-google-adk/google-mcp-security-agent/agent.py b/run-with-google-adk/google-mcp-security-agent/agent.py index 10b2e838..b53f162c 100644 --- a/run-with-google-adk/google-mcp-security-agent/agent.py +++ b/run-with-google-adk/google-mcp-security-agent/agent.py @@ -34,6 +34,8 @@ async def get_all_tools(): gti_tools = [] secops_soar_tools = [] scc_tools = [] # Initialize scc_tools + compliance_tools = [] # Initialize compliance_tools + exit_stack = AsyncExitStack() uv_dir_prefix="../server" env_file_path = "../../../run-with-google-adk/google-mcp-security-agent/.env" @@ -98,11 +100,23 @@ async def get_all_tools(): "scc_mcp.py" ], ),async_exit_stack=exit_stack - ) + ) + + if os.environ.get("LOAD_COMPLIANCE_MCP") == "Y": + compliance_tools, exit_stack = await MCPToolset.from_server( + connection_params=StdioServerParameters( + command='uv', + args=[ "--directory", + uv_dir_prefix + "/compliance/", + "run", + "compliance_mcp.py" + ], + ),async_exit_stack=exit_stack + ) logging.info("MCP Toolsets created successfully.") - return secops_tools+gti_tools+secops_soar_tools+scc_tools, exit_stack + return secops_tools+gti_tools+secops_soar_tools+scc_tools+compliance_tools, exit_stack def make_tools_gemini_compatible(tools): """ diff --git a/server/compliance/README.md b/server/compliance/README.md index cc297229..26093f9d 100644 --- a/server/compliance/README.md +++ b/server/compliance/README.md @@ -12,6 +12,17 @@ This is an MCP (Model Context Protocol) server for interacting with Google Cloud - `parent` (required): The parent resource name (e.g., 'organizations/123456', 'folders/123456', or 'projects/my-project'). - `page_size` (optional): The maximum number of frameworks to return. Defaults to 50. +- **`list_constraints(parent)`** + - **Description**: Returns details of all available Org Policy constraints for a given resource: (organization, folder, or project). + - **Parameters**: + - `parent` (required): The parent resource name (e.g., 'organizations/123456', 'folders/123456', or 'projects/my-project'). + +- **`list_active_policies(parent)`** + - **Description**: Returns details of all active/enforced Org Policy policies for a given resource: (organization, folder, or project). + - **Parameters**: + - `parent` (required): The parent resource name (e.g., 'organizations/123456', 'folders/123456', or 'projects/my-project'). + + ## Configuration ### MCP Server Configuration @@ -68,4 +79,22 @@ await list_frameworks("organizations/123456789") # List frameworks for a project with custom page size await list_frameworks("projects/my-project-id", page_size=25) +``` + +### List Constraints +```python +# List constraints for an organization +await list_constraints("organizations/123456789") + +# List constraints for a project +await list_constraints("projects/my-project-id") +``` + +### List Active Org Policies +```python +# List constraints for an organization +await list_active_policies("organizations/123456789") + +# List constraints for a project +await list_active_policies("projects/my-project-id") ``` \ No newline at end of file diff --git a/server/compliance/compliance_mcp.py b/server/compliance/compliance_mcp.py index de672e96..2fde688f 100644 --- a/server/compliance/compliance_mcp.py +++ b/server/compliance/compliance_mcp.py @@ -15,6 +15,7 @@ import os from typing import Any, Dict, List +from google.cloud import orgpolicy_v2 from google.api_core import exceptions as google_exceptions from google.cloud.cloudsecuritycompliance_v1alpha.services.config import ConfigClient from google.protobuf import json_format @@ -46,6 +47,17 @@ logger.error(f"Failed to initialize Cloud Security Compliance Config Client: {e}", exc_info=True) config_client = None +# --- Org Policy Client Initialization --- +# The client automatically uses Application Default Credentials (ADC). +# Ensure ADC are configured in the environment where the server runs +# (e.g., by running `gcloud auth application-default login`). +try: + orgpolicy_client = orgpolicy_v2.OrgPolicyClient() + logger.info("Successfully initialized Org Policy Client.") +except Exception as e: + logger.error(f"Failed to initialize Org Policy Client: {e}", exc_info=True) + orgpolicy_client = None # Indicate client is not available + # --- Helper Function for Proto to Dict Conversion --- def proto_message_to_dict(message: Any) -> Dict[str, Any]: @@ -60,6 +72,7 @@ def proto_message_to_dict(message: Any) -> Dict[str, Any]: # --- Cloud Security Compliance Tools --- +# --- List Frameworks tool --- @mcp.tool() async def list_frameworks( parent: str, @@ -70,7 +83,7 @@ async def list_frameworks( Description: Lists available compliance frameworks for a given parent resource (organization, folder, or project). Returns information about available compliance frameworks and their details. Parameters: - parent (required): The parent resource name (e.g., 'organizations/123456', 'folders/123456', or 'projects/my-project'). + parent (required): The parent resource name. Currently only supports 'organization' type resource (e.g., 'organizations/123456'). page_size (optional): The maximum number of frameworks to return. Defaults to 50. """ if not config_client: @@ -79,10 +92,13 @@ async def list_frameworks( logger.info(f"Listing frameworks for parent: {parent}") try: + # Append "/locations/global" to parent as required by the API request_args = { - "parent": parent, + "parent": f"{parent}/locations/global", "page_size": page_size, } + logger.info(f"Request args for list_frameworks: {request_args}") + response_pager = config_client.list_frameworks(request=request_args) @@ -127,6 +143,175 @@ async def list_frameworks( logger.error(f"An unexpected error occurred listing frameworks: {e}", exc_info=True) return {"error": "An unexpected error occurred", "details": str(e)} + +# --- Describe Framework tool --- +@mcp.tool() +async def describe_framework( + parent: str, + framework_name: str, +) -> Dict[str, Any]: + """Name: describe_framework + + Description: Returns the list of control descriptions under the given framework for a parent resource. + Parameters: + parent (required): The parent resource name. Currently only supports 'organization' type resource (e.g., 'organizations/123456'). + framework_name (required): The full resource name of the framework to describe. + """ + if not config_client: + return {"error": "Cloud Security Compliance Config Client not initialized."} + + logger.info(f"Describing framework '{framework_name}' for parent: {parent}") + + try: + # Step 1: List all frameworks for the parent + request_args = { + "parent": f"{parent}/locations/global", + "page_size": 100, + } + all_frameworks = config_client.list_frameworks(request=request_args) + selected_framework = None + + # Step 2: Filter out the framework matching the given name + + for framework in all_frameworks.frameworks: + if framework.name == framework_name: + selected_framework = framework + break + + logger.info(f"Selected framework: {selected_framework}") + + if not selected_framework: + logger.error(f"Framework '{framework_name}' not found under parent '{parent}'.") + return {"error": "Framework not found", "details": f"Framework '{framework_name}' not found under parent '{parent}'."} + + # Step 3: Get control names from the selected framework + control_names = [control.name for control in getattr(selected_framework, "cloud_control_details", [])] + if not control_names: + return {"controls": []} + + # Step 4: Call ListCloudControls API to get control descriptions + controls_request = { + "parent": f"{parent}/locations/global", + } + controls_response = config_client.list_cloud_controls(request=controls_request) + + # Step 5: Return a list of cloud control descriptions + control_descriptions = [] + for control in controls_response.cloud_controls: + if control.name in control_names: + control_descriptions.append(control.description) + + return { + "controls": control_descriptions + } + + except google_exceptions.NotFound as e: + logger.error(f"Parent resource or framework not found: {parent}: {e}") + return {"error": "Not Found", "details": f"Could not find parent resource or framework. {str(e)}"} + except google_exceptions.PermissionDenied as e: + logger.error(f"Permission denied for describing framework on {parent}: {e}") + return {"error": "Permission Denied", "details": str(e)} + except google_exceptions.InvalidArgument as e: + logger.error(f"Invalid argument for describing framework on {parent}: {e}") + return {"error": "Invalid Argument", "details": str(e)} + except Exception as e: + logger.error(f"An unexpected error occurred describing framework: {e}", exc_info=True) + return {"error": "An unexpected error occurred", "details": str(e)} + +# --- List Constraints tool --- +@mcp.tool() +async def list_constraints( + parent: str +) -> Dict[str, str]: + """Name: list all available constraints + Description: Returns details of all available Org Policy constraints for a given resource:[organization/project/folder]. + Parameters: + parent (required): The Google Cloud resource that parents the constraint. Must be in one of the following forms: + * `projects/{project_number}` + * `projects/{project_id}` + * `folders/{folder_id}` + * `organizations/{organization_id}` + """ + if not orgpolicy_client: + return {"error": "Org Polciy Client not initialized."} + + logger.info(f"getting all constraints for parent: {parent}") + + try: + + constraints_iter = orgpolicy_client.list_constraints(request={"parent": parent}) + constraints_map = {} + + for constraint in constraints_iter: + # Each constraint has .name, .display_name, .description, etc. + constraints_map[constraint.name] = { + "display_name": constraint.display_name, + "description": constraint.description + } + + return constraints_map + + except Exception as e: + logger.error(f"Error listing constraints for {parent}: {e}", exc_info=True) + return {"error": str(e)} + + +# --- List Active Policies tool --- +@mcp.tool() +async def list_active_policies( + parent: str +) -> Dict[str, Any]: + """Name: list active/enforced policy details + Description: Returns details of all active/enforced Org Policy policies for a given resource:[organization/project/folder]. + Parameters: + parent (required): The Google Cloud resource that parents the policies. Must be in one of the following forms: + * `projects/{project_number}` + * `projects/{project_id}` + * `folders/{folder_id}` + * `organizations/{organization_id}` + """ + if not orgpolicy_client: + return {"error": "Org Polciy Client not initialized."} + + logger.info(f"getting all active policies for parent: {parent}") + + try: + # Step 1: Fetch all policies and filter for enforced ones + policies = orgpolicy_client.list_policies(request={"parent": parent}) + active_policies_constraint_id_set = set() + + for policy in policies: + rules = getattr(policy.spec, "rules", []) if policy.spec else [] + is_enforced = any(getattr(rule, "enforce", False) for rule in rules) + if is_enforced: + policy_name = policy.name + # Extract constraint_id from policy_name + # Example: "projects/{project_number}/policies/{constraint_id}" + if "/policies/" in policy_name: + constraint_id = policy_name.split("/policies/")[-1] + active_policies_constraint_id_set.add(constraint_id) + else: + continue + + # Step 2: Fetch all constraints for the parent + constraints = orgpolicy_client.list_constraints(request={"parent": parent}) + constraint_desc_map = {} + for constraint in constraints: + # constraint.name is like "projects/123/constraints/serviceuser.services" + # We want to map by just the constraint_name part + constraint_id = constraint.name.split("/constraints/")[-1] + if constraint_id in active_policies_constraint_id_set: + constraint_desc_map[constraint_id] = { + "description": constraint.description, + "display_name": constraint.display_name} + + return constraint_desc_map + + except Exception as e: + logger.error(f"Error listing active policies for {parent}: {e}", exc_info=True) + return {"error": str(e)} + + # --- Main execution --- def main() -> None: diff --git a/server/compliance/setup.py b/server/compliance/setup.py index dd2b1be1..82394080 100644 --- a/server/compliance/setup.py +++ b/server/compliance/setup.py @@ -23,6 +23,7 @@ install_requires=[ "mcp", "google-cloud-cloudsecuritycompliance", + "google-cloud-org-policy", ], entry_points={ "console_scripts": [