awslabs.well-architected-security-mcp-server 0.1.1__py3-none-any.whl

awslabs/well_architected_security_mcp_server/util/prompt_utils.py

@@ -0,0 +1,173 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# 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.
+
+"""Utilities for working with prompt templates."""
+
+import os
+import re
+from typing import Any, Dict, List, Optional
+
+from loguru import logger
+
+# Cache for prompt templates
+_prompt_templates = {}
+_is_initialized = False
+
+
+def load_prompt_templates(file_path: str = "PROMPT_TEMPLATE.md") -> Dict[str, Dict[str, Any]]:
+    """Load prompt templates from a markdown file.
+
+    Args:
+        file_path: Path to the markdown file containing prompt templates
+
+    Returns:
+        Dictionary mapping template names to template content and metadata
+    """
+    global _prompt_templates, _is_initialized
+
+    if _is_initialized:
+        return _prompt_templates
+
+    try:
+        # Check if file exists
+        if not os.path.exists(file_path):
+            logger.error(f"Prompt template file not found: {file_path}")
+            return {}
+
+        # Read the file content
+        with open(file_path, "r") as f:
+            content = f.read()
+
+        # Parse the markdown content to extract templates
+        # First, split by level 2 headers (## )
+        sections = re.split(r"(?m)^## ", content)
+
+        # The first section is the introduction, skip it
+        if sections and not sections[0].startswith("## "):
+            sections = sections[1:]
+
+        # Process each section
+        for section in sections:
+            if not section.strip():
+                continue
+
+            # Extract the section title (template name)
+            lines = section.split("\n")
+            title = lines[0].strip()
+
+            # Convert title to a template name (lowercase, underscores)
+            template_name = title.lower().replace(" ", "_")
+
+            # Extract the template content (between triple backticks)
+            template_content = ""
+            description = ""
+            in_code_block = False
+            desc_lines = []
+
+            for line in lines[1:]:
+                if line.strip() == "```":
+                    in_code_block = not in_code_block
+                    continue
+
+                if in_code_block:
+                    template_content += line + "\n"
+                elif line.strip() and not in_code_block:
+                    desc_lines.append(line.strip())
+
+            # Join description lines
+            if desc_lines:
+                description = " ".join(desc_lines)
+
+            # Store the template
+            _prompt_templates[template_name] = {
+                "name": template_name,
+                "title": title,
+                "content": template_content.strip(),
+                "description": description,
+            }
+
+        # Special handling for the workflow example section
+        workflow_section = None
+        for section in content.split("## "):
+            if section.startswith("Example Workflow"):
+                workflow_section = section
+                break
+
+        if workflow_section:
+            _prompt_templates["workflow_example"] = {
+                "name": "workflow_example",
+                "title": "Example Workflow",
+                "content": workflow_section.split("Example Workflow")[1].strip(),
+                "description": "Recommended workflow for a comprehensive security assessment",
+            }
+
+        _is_initialized = True
+        logger.info(f"Loaded {len(_prompt_templates)} prompt templates from {file_path}")
+        return _prompt_templates
+
+    except Exception as e:
+        logger.error(f"Error loading prompt templates: {e}")
+        return {}
+
+
+def get_prompt_template(template_name: str) -> Optional[Dict[str, Any]]:
+    """Get a specific prompt template by name.
+
+    Args:
+        template_name: Name of the template to retrieve
+
+    Returns:
+        Template dictionary or None if not found
+    """
+    global _prompt_templates, _is_initialized
+
+    if not _is_initialized:
+        load_prompt_templates()
+
+    return _prompt_templates.get(template_name)
+
+
+def get_all_template_names() -> List[str]:
+    """Get a list of all available template names.
+
+    Returns:
+        List of template names
+    """
+    global _prompt_templates, _is_initialized
+
+    if not _is_initialized:
+        load_prompt_templates()
+
+    return list(_prompt_templates.keys())
+
+
+def get_template_metadata() -> List[Dict[str, str]]:
+    """Get metadata for all available templates.
+
+    Returns:
+        List of dictionaries with template metadata
+    """
+    global _prompt_templates, _is_initialized
+
+    if not _is_initialized:
+        load_prompt_templates()
+
+    return [
+        {
+            "name": template["name"],
+            "title": template["title"],
+            "description": template["description"],
+        }
+        for template in _prompt_templates.values()
+    ]

awslabs/well_architected_security_mcp_server/util/resource_utils.py

@@ -0,0 +1,109 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# 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.
+
+"""General utility functions for AWS resource operations."""
+
+from typing import Any, Dict
+
+import boto3
+from botocore.config import Config
+from mcp.server.fastmcp import Context
+
+from awslabs.well_architected_security_mcp_server import __version__
+
+# User agent configuration for AWS API calls
+USER_AGENT_CONFIG = Config(
+    user_agent_extra=f"awslabs/mcp/well-architected-security-mcp-server/{__version__}"
+)
+
+
+async def list_services_in_region(
+    region: str, session: boto3.Session, ctx: Context
+) -> Dict[str, Any]:
+    """List all AWS services being used in a specific region.
+
+    Args:
+        region: AWS region to list services for
+        session: boto3 Session for AWS API calls
+        ctx: MCP context for error reporting
+
+    Returns:
+        Dictionary with services information and counts
+    """
+    try:
+        # Initialize the result dictionary
+        result = {"region": region, "services": [], "service_counts": {}, "total_resources": 0}
+
+        # Use Resource Explorer to efficiently discover resources
+        try:
+            resource_explorer = session.client(
+                "resource-explorer-2", region_name=region, config=USER_AGENT_CONFIG
+            )
+
+            # Check if Resource Explorer is available in this region
+            try:
+                # Try to search with Resource Explorer
+                resource_explorer.search(
+                    QueryString="*",
+                    MaxResults=1,  # Just checking if it works
+                )
+            except Exception as e:
+                if "Resource Explorer has not been set up" in str(e):
+                    await ctx.warning(
+                        f"Resource Explorer not set up in {region}. Using alternative method."
+                    )
+                    return {"region": region, "services": [], "error": str(e)}
+                else:
+                    raise e
+
+            # Resource Explorer is available, use it to get all resources
+            paginator = resource_explorer.get_paginator("search")
+            page_iterator = paginator.paginate(QueryString="*", MaxResults=1000)
+
+            # Track unique services
+            services_set = set()
+            service_resource_counts = {}
+
+            # Process each page of results
+            for page in page_iterator:
+                for resource in page.get("Resources", []):
+                    # Extract service from ARN
+                    arn = resource.get("Arn", "")
+                    if arn:
+                        arn_parts = arn.split(":")
+                        if len(arn_parts) >= 3:
+                            service = arn_parts[2]
+                            services_set.add(service)
+
+                            # Update count for this service
+                            if service in service_resource_counts:
+                                service_resource_counts[service] += 1
+                            else:
+                                service_resource_counts[service] = 1
+
+            # Update result with discovered services
+            result["services"] = sorted(list(services_set))
+            result["service_counts"] = service_resource_counts
+            result["total_resources"] = sum(service_resource_counts.values())
+
+        except Exception as e:
+            await ctx.warning(f"Error using Resource Explorer in {region}: {e}")
+            # Fall back to alternative method
+            return {"region": region, "services": [], "error": str(e)}
+
+        return result
+
+    except Exception as e:
+        await ctx.error(f"Error listing services in region {region}: {e}")
+        return {"region": region, "services": [], "error": str(e)}