fix(diagrams): adds ast parsing to block import and from import

src/aws-diagram-mcp-server/awslabs/aws_diagram_mcp_server/scanner.py

@@ -22,6 +22,11 @@
 
 # Suppress AST deprecation warnings for Python 3.14 compatibility
 warnings.filterwarnings('ignore', category=DeprecationWarning, module='ast')
+# Suppress deprecation warnings from bandit and other libraries using deprecated AST features
+warnings.filterwarnings('ignore', category=DeprecationWarning, message=r'.*ast\.Bytes.*')
+warnings.filterwarnings(
+    'ignore', category=DeprecationWarning, message='.*Attribute n is deprecated.*'
+)
 
 
 class SecurityIssue(BaseModel):
@@ -57,7 +62,15 @@ class CodeScanResult(BaseModel):
 async def validate_syntax(code: str) -> Tuple[bool, Optional[str]]:
     """Validate Python code syntax using ast."""
     try:
-        ast.parse(code)
+        tree = ast.parse(code)
+
+        # Check for import statements
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Import):
+                return False, f'Import statements are not allowed (line {node.lineno})'
+            elif isinstance(node, ast.ImportFrom):
+                return False, f'Import statements are not allowed (line {node.lineno})'
+
         return True, None
     except SyntaxError as e:
         error_msg = f'Syntax error at line {e.lineno}: {e.msg}'
@@ -225,6 +238,7 @@ def check_dangerous_functions(code: str) -> List[Dict[str, Any]]:
         'os.popen',
         '__import__',
         'pickle.loads',
+        'spawn(',
     ]
 
     results = []

src/aws-diagram-mcp-server/pyproject.toml

@@ -3,33 +3,30 @@ name = "awslabs.aws-diagram-mcp-server"
 version = "1.0.10"
 description = "An MCP server that seamlessly creates diagrams using the Python diagrams package DSL"
 readme = "README.md"
-requires-python = ">=3.10"
+requires-python = ">=3.12"
 dependencies = [
+    "bandit>=1.8.6",
+    "boto3>=1.40.53",
     "diagrams>=0.24.4",
-    "boto3>=1.37.27",
-    "mcp[cli]>=1.11.0",
-    "pydantic>=2.10.6",
-    "bandit>=1.7.5",
-    "sarif-om>=1.0.0",     # Fixes GitHub issue #1041: No module named 'sarif_om'
-    # Security fixes for CVEs
-    "setuptools>=78.1.1",  # Fixes CVE-2025-47273
-    "starlette>=0.47.2",   # Fixes CVE-2025-54121
-    "urllib3>=2.5.0",      # Fixes CVE-2025-50181, CVE-2025-50182
+    "mcp[cli]>=1.17.0",
+    "pydantic>=2.12.2",
+    "sarif-om>=1.0.4",
+    "setuptools>=80.9.0",
+    "starlette>=0.48.0",
+    "urllib3>=2.5.0",
 ]
-license = {text = "Apache-2.0"}
-license-files = ["LICENSE", "NOTICE" ]
+license = { text = "Apache-2.0" }
+license-files = ["LICENSE", "NOTICE"]
 
 authors = [
-    {name = "Amazon Web Services"},
-    {name = "AWSLabs MCP", email="[email protected]"},
+    { name = "Amazon Web Services" },
+    { name = "AWSLabs MCP", email = "[email protected]" },
 ]
 classifiers = [
     "License :: OSI Approved :: Apache Software License",
     "Operating System :: OS Independent",
     "Programming Language :: Python",
     "Programming Language :: Python :: 3",
-    "Programming Language :: Python :: 3.10",
-    "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
 ]
@@ -46,13 +43,13 @@ Changelog = "https://github.com/awslabs/mcp/blob/main/src/aws-diagram-mcp-server
 
 [dependency-groups]
 dev = [
-    "commitizen>=4.2.2",
-    "pre-commit>=4.1.0",
-    "ruff>=0.9.7",
-    "pyright>=1.1.398",
-    "pytest>=7.4.0",
-    "pytest-cov>=4.1.0",
-    "pytest-asyncio>=0.26.0",
+    "commitizen>=4.9.1",
+    "pre-commit>=4.2.0",
+    "pyright>=1.1.406",
+    "pytest>=8.4.2",
+    "pytest-asyncio>=1.2.0",
+    "pytest-cov>=7.0.0",
+    "ruff>=0.14.1",
 ]
 
 [build-system]
@@ -74,7 +71,7 @@ exclude = [
     "**/env",
     "**/.ruff_cache",
     "**/.venv",
-    "**/.ipynb_checkpoints"
+    "**/.ipynb_checkpoints",
 ]
 force-exclude = true
 
@@ -102,7 +99,13 @@ docstring-code-format = true
 
 [tool.pyright]
 include = ["awslabs", "tests"]
-exclude = ["**/__pycache__", "**/.venv", "**/node_modules", "**/dist", "**/build"]
+exclude = [
+    "**/__pycache__",
+    "**/.venv",
+    "**/node_modules",
+    "**/dist",
+    "**/build",
+]
 typeCheckingMode = "basic"
 reportMissingImports = false
 reportUnusedExpression = false
@@ -120,7 +123,7 @@ version = "0.0.0"
 tag_format = "v$version"
 version_files = [
     "pyproject.toml:version",
-    "awslabs/aws_diagram_mcp_server/__init__.py:__version__"
+    "awslabs/aws_diagram_mcp_server/__init__.py:__version__",
 ]
 update_changelog_on_bump = true
 
@@ -130,7 +133,7 @@ packages = ["awslabs"]
 [tool.bandit]
 # Skip specific issues
 skips = ["B102"]
-exclude_dirs = ["venv","tests"]
+exclude_dirs = ["venv", "tests"]
 
 # Per-file skips
 per_file_skips = { "awslabs/aws_diagram_mcp_server/diagrams.py" = ["B102"] }
@@ -139,9 +142,7 @@ per_file_skips = { "awslabs/aws_diagram_mcp_server/diagrams.py" = ["B102"] }
 testpaths = "tests"
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
-markers = [
-    "asyncio: mark a test as an asyncio coroutine",
-]
+markers = ["asyncio: mark a test as an asyncio coroutine"]
 filterwarnings = [
     "ignore::DeprecationWarning:ast",
     "ignore:ast.Str is deprecated:DeprecationWarning",

src/aws-diagram-mcp-server/tests/conftest.py

@@ -16,10 +16,29 @@
 
 import pytest
 import tempfile
+import warnings
 from awslabs.aws_diagram_mcp_server.models import DiagramType
 from typing import Dict, Generator
 
 
+# Suppress AST deprecation warnings from bandit and other libraries
+warnings.filterwarnings('ignore', category=DeprecationWarning, message=r'.*ast\.Bytes.*')
+warnings.filterwarnings(
+    'ignore', category=DeprecationWarning, message='.*Attribute n is deprecated.*'
+)
+
+
+@pytest.fixture(autouse=True)
+def suppress_deprecation_warnings():
+    """Suppress deprecation warnings for all tests."""
+    with warnings.catch_warnings():
+        warnings.filterwarnings('ignore', category=DeprecationWarning, message=r'.*ast\.Bytes.*')
+        warnings.filterwarnings(
+            'ignore', category=DeprecationWarning, message='.*Attribute n is deprecated.*'
+        )
+        yield
+
+
 @pytest.fixture
 def temp_workspace_dir() -> Generator[str, None, None]:
     """Create a temporary directory for diagram output."""

src/aws-diagram-mcp-server/tests/test_sarif_fix.py

@@ -61,17 +61,17 @@ def test_bandit_security_scan(self):
         import asyncio
         from awslabs.aws_diagram_mcp_server.scanner import scan_python_code
 
-        # Test with code that should trigger security warnings
+        # Test with code that should trigger security warnings (without imports)
         unsafe_code = """
-import subprocess
-subprocess.call("echo test", shell=True)
+exec("print('Hello')")
+eval("2 + 2")
 """
 
         async def run_test():
             result = await scan_python_code(unsafe_code)
             assert result is not None
             assert result.syntax_valid is True
-            # Should have security issues due to subprocess.call with shell=True
+            # Should have security issues due to exec/eval
             assert result.has_errors is True
             assert len(result.security_issues) > 0
             return result
@@ -80,11 +80,11 @@ async def run_test():
         result = asyncio.run(run_test())
 
         # Verify we found the expected security issue
-        found_subprocess_issue = any(
-            'subprocess' in issue.issue_text.lower() or 'shell' in issue.issue_text.lower()
+        found_dangerous_function = any(
+            'exec' in issue.issue_text.lower() or 'eval' in issue.issue_text.lower()
             for issue in result.security_issues
         )
-        assert found_subprocess_issue, 'Should detect subprocess security issue'
+        assert found_dangerous_function, 'Should detect exec/eval security issue'
 
     def test_diagram_generation_unaffected(self):
         """Test that diagram generation still works normally after the fix."""
@@ -96,8 +96,6 @@ def test_diagram_generation_unaffected(self):
         # Simple diagram code
         diagram_code = """
 with Diagram("Test Diagram", show=False):
-    from diagrams.aws.compute import EC2
-    from diagrams.aws.database import RDS
 
     web = EC2("Web Server")
     db = RDS("Database")

src/aws-diagram-mcp-server/tests/test_scanner.py

@@ -61,6 +61,30 @@ def factorial(n):
         assert valid is True
         assert error is None
 
+    @pytest.mark.asyncio
+    async def test_import_rejected(self):
+        """Test that import statements are rejected."""
+        code = """
+import os
+print("Hello")
+"""
+        valid, error = await validate_syntax(code)
+        assert valid is False
+        assert error is not None
+        assert 'Import statements are not allowed' in error
+
+    @pytest.mark.asyncio
+    async def test_import_from_rejected(self):
+        """Test that from...import statements are rejected."""
+        code = """
+from typing import List
+print("Hello")
+"""
+        valid, error = await validate_syntax(code)
+        assert valid is False
+        assert error is not None
+        assert 'Import statements are not allowed' in error
+
 
 class TestSecurityChecking:
     """Tests for the security checking functionality."""
@@ -248,8 +272,8 @@ def add(a, b):
     async def test_security_issue(self):
         """Test scanning code with security issues."""
         code = """
-import os
-os.system("rm -rf /")  # This is dangerous
+exec("malicious_code")
+eval("dangerous_expression")
 """
         result = await scan_python_code(code)
         assert result.has_errors is True