Add the ability to use Cloudformation hosted Tranforms

src/cloud_radar/cf/unit/_template.py

@@ -3,7 +3,7 @@
 import json
 import re
 from pathlib import Path
-from typing import Any, Callable, Dict, Generator, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
 
 import yaml  # noqa: I100
 from cfn_tools import dump_yaml, load_yaml  # type: ignore  # noqa: I100, I201
@@ -75,6 +75,9 @@
         self.Region = Template.Region
         self.imports = imports
         self.dynamic_references = dynamic_references
+        self.transforms: Optional[Union[str, List[str]]] = self.template.get(
+            "Transform", None
+        )
 
     @classmethod
     def from_yaml(
@@ -197,12 +200,54 @@
             # If we get this far then we do not support this type of configuration file
             raise ValueError("Parameter file is not in a supported format")
 
+    def load_allowed_functions(self) -> functions.Dispatch:
+        """Loads the allowed functions for this template.
+
+        Raises:
+            ValueError: If the transform is not supported.
+            ValueError: If the Transform section is not a string or list.
+
+        Returns:
+            functions.Dispatch: A dictionary of allowed functions.
+        """
+        if self.transforms is None:
+            return functions.ALL_FUNCTIONS
+
+        if isinstance(self.transforms, str):
+            if self.transforms not in functions.TRANSFORMS:
+                raise ValueError(f"Transform {self.transforms} not supported")
+
+            # dict of transform functions
+            transform_functions = functions.TRANSFORMS[self.transforms]
+
+            # return the merger of ALL_FUNCTIONS and the transform functions
+            return {**functions.ALL_FUNCTIONS, **transform_functions}
+
+        if isinstance(self.transforms, list):
+            # dict of transform functions
+            transform_functions = {}
+            for transform in self.transforms:
+                if transform not in functions.TRANSFORMS:
+                    raise ValueError(f"Transform {transform} not supported")
+                transform_functions = {
+                    **transform_functions,
+                    **functions.TRANSFORMS[transform],
+                }
+
+            # return the merger of ALL_FUNCTIONS and the transform functions
+            return {**functions.ALL_FUNCTIONS, **transform_functions}
+
+        raise ValueError(f"Transform {self.transforms} not supported")
+
     def render_all_sections(self, template: Dict[str, Any]) -> Dict[str, Any]:
         """Solves all conditionals, references and pseudo variables for all sections"""
+
+        allowed_functions = self.load_allowed_functions()
+
         if "Conditions" in template:
             template["Conditions"] = self.resolve_values(
                 template["Conditions"],
-                functions.ALL_FUNCTIONS,
+                allowed_functions,
             )
 
         template_sections = ["Resources", "Outputs"]
@@ -228,7 +273,7 @@
 
                 template[section][r_name] = self.resolve_values(
                     r_value,
-                    functions.ALL_FUNCTIONS,
+                    allowed_functions,
                 )
 
         return template

src/cloud_radar/cf/unit/functions.py

@@ -15,8 +15,16 @@
 if TYPE_CHECKING:
     from ._template import Template
 
+# Dispatch represents a dictionary where the keys are Cloudformation
+# function names in there long form and the values are the functions
+# in python that solve them.
 Dispatch = Dict[str, Callable[..., Any]]
 
+# Mapping represents the Cloudformation Mappings section of a template.
+# The keys are the names of the maps and the values are the maps themselves.
+# The maps are a nested dictionary.
+Mapping = Dict[str, Dict[str, Dict[str, Any]]]
+
 REGION_DATA = None
 
 
@@ -281,6 +289,39 @@
     return condition_value
 
 
+def _find_in_map(maps: Mapping, map_name: str, top_key: str, second_key: str) -> Any:
+    """Solves AWS FindInMap intrinsic function.
+
+    Args:
+        maps (Mapping): The Cloudformation Mappings section of the template.
+        map_name (str): The name of the Map to search.
+        top_key (str): The top level key to search.
+        second_key (str): The second level key to search.
+
+    Raises:
+        KeyError: If map_name is not found in the Mapping section.
+        KeyError: If top_key is not found in the Map.
+        KeyError: If second_key is not found in the Map.
+
+    Returns:
+        Any: The requested value from the Map.
+    """
+    if map_name not in maps:
+        raise KeyError(f"Unable to find {map_name} in Mappings section of template.")
+
+    map = maps[map_name]
+
+    if top_key not in map:
+        raise KeyError(f"Unable to find key {top_key} in map {map_name}.")
+
+    first_level = map[top_key]
+
+    if second_key not in first_level:
+        raise KeyError(f"Unable to find key {second_key} in map {map_name}.")
+
+    return first_level[second_key]
+
+
 def find_in_map(template: "Template", values: Any) -> Any:
     """Solves AWS FindInMap intrinsic function.
 
@@ -319,20 +360,57 @@
 
     maps = template.template["Mappings"]
 
-    if map_name not in maps:
-        raise KeyError(f"Unable to find {map_name} in Mappings section of template.")
+    return _find_in_map(maps, map_name, top_key, second_key)
 
-    map = maps[map_name]
 
-    if top_key not in map:
-        raise KeyError(f"Unable to find key {top_key} in map {map_name}.")
+def enhanced_find_in_map(template: "Template", values: Any) -> Any:
+    """Solves AWS FindInMap intrinsic function. This version allows for a default value.
 
-    first_level = map[top_key]
+    Args:
+        template (Template): The template being tested.
+        values (Any): The values passed to the function.
 
-    if second_key not in first_level:
-        raise KeyError(f"Unable to find key {second_key} in map {map_name}.")
+    Raises:
+        TypeError: If values is not a list.
+        ValueError: If length of values is not 3.
+        KeyError: If the Map or specified keys are missing.
 
-    return first_level[second_key]
+    Returns:
+        Any: The requested value from the Map.
+    """
+
+    if not isinstance(values, list):
+        raise TypeError(
+            f"Fn::FindInMap - The values must be a List, not {type(values).__name__}."
+        )
+
+    if len(values) not in [3, 4]:
+        raise ValueError(
+            (
+                "Fn::FindInMap - The values must contain "
+                "a MapName, TopLevelKey and SecondLevelKey. "
+                "Optionally, a third value can be provided to "
+                "specify a default value."
+            )
+        )
+
+    map_name = values[0]
+    top_key = values[1]
+    second_key = values[2]
+
+    if "Mappings" not in template.template:
+        raise KeyError("Unable to find Mappings section in template.")
+
+    maps = template.template["Mappings"]
+
+    default_value: Dict[str, Any] = values.pop(3) if len(values) == 4 else {}
+
+    try:
+        return _find_in_map(maps, map_name, top_key, second_key)
+    except KeyError:
+        if "DefaultValue" in default_value:
+            return default_value["DefaultValue"]
+        raise
 
 
 def get_att(template: "Template", values: Any) -> str:
@@ -941,3 +1019,15 @@
     "Fn::Transform": {},  # Transform isn't fully implemented
     "Ref": {},  # String only.
 }
+
+# Extra functions that are allowed if the template is using a transform.
+TRANSFORMS: Dict[str, Dispatch] = {
+    "AWS::CodeDeployBlueGreen": {},
+    "AWS::Include": {},
+    "AWS::LanguageExtensions": {
+        "Fn::FindInMap": enhanced_find_in_map,
+    },
+    "AWS::SecretsManager-2020-07-23": {},
+    "AWS::Serverless-2016-10-31": {},
+    "AWS::ServiceCatalog": {},
+}

src/cloud_radar/cf/unit/test__template.py

@@ -0,0 +1,46 @@
+# Just a note that GitHub Copilot generated this entire file, first try
+import pytest
+
+from cloud_radar.cf.unit import functions
+from cloud_radar.cf.unit._template import Template
+
+
+def test_load_allowed_functions_no_transforms():
+    template = Template({})
+    result = template.load_allowed_functions()
+    assert result == functions.ALL_FUNCTIONS
+
+
+def test_load_allowed_functions_single_transform():
+    template = Template({"Transform": "AWS::Serverless-2016-10-31"})
+    result = template.load_allowed_functions()
+    expected = {
+        **functions.ALL_FUNCTIONS,
+        **functions.TRANSFORMS["AWS::Serverless-2016-10-31"],
+    }
+    assert result == expected
+
+
+def test_load_allowed_functions_multiple_transforms():
+    template = Template({"Transform": ["AWS::Serverless-2016-10-31", "AWS::Include"]})
+    result = template.load_allowed_functions()
+    expected = {
+        **functions.ALL_FUNCTIONS,
+        **functions.TRANSFORMS["AWS::Serverless-2016-10-31"],
+        **functions.TRANSFORMS["AWS::Include"],
+    }
+    assert result == expected
+
+
+def test_load_allowed_functions_invalid_transform():
+    template = Template({"Transform": "InvalidTransform"})
+    with pytest.raises(ValueError):
+        template.load_allowed_functions()
+
+
+def test_load_allowed_functions_invalid_transforms():
+    template = Template(
+        {"Transform": ["AWS::Serverless-2016-10-31", "InvalidTransform"]}
+    )
+    with pytest.raises(ValueError):
+        template.load_allowed_functions()

tests/templates/test_maps.yml

@@ -0,0 +1,20 @@
+AWSTemplateFormatVersion: 2010-09-09
+Description: "Creates an S3 bucket to store logs."
+
+Transform: AWS::LanguageExtensions
+
+Mappings:
+  Test:
+    Foo:
+      bar: baz
+
+Resources:
+  BazBucket:
+    Type: AWS::S3::Bucket
+    Properties:
+      BucketName: !FindInMap [Test, Foo, bar]
+
+  BazingaBucket:
+    Type: AWS::S3::Bucket
+    Properties:
+      BucketName: !FindInMap [Test, Foo, baz, DefaultValue: bazinga]

tests/test_cf/test_examples/test_unit.py

@@ -12,6 +12,13 @@ def template():
     return Template.from_yaml(template_path.resolve(), {})
 
 
+@pytest.fixture
+def map_template():
+    template_path = Path(__file__).parent / "../../templates/test_maps.yml"
+
+    return Template.from_yaml(template_path.resolve(), {})
+
+
 def test_log_defaults(template: Template):
     stack = template.create_stack({"BucketPrefix": "testing"})
 
@@ -46,3 +53,15 @@ def test_log_retain(template: Template):
     always_true = stack.get_condition("AlwaysTrue")
 
     always_true.assert_value_is(True)
+
+
+def test_maps(map_template: Template):
+    stack = map_template.create_stack()
+
+    baz_bucket = stack.get_resource("BazBucket")
+
+    baz_bucket.assert_property_has_value("BucketName", "baz")
+
+    bazinga_bucket = stack.get_resource("BazingaBucket")
+
+    bazinga_bucket.assert_property_has_value("BucketName", "bazinga")

tests/test_cf/test_unit/test_functions.py

@@ -278,6 +278,33 @@ def test_find_in_map():
     assert result == expected
 
 
+def test_transform_find_in_map():
+    template = {}
+
+    add_metadata(template, Template.Region)
+
+    template = Template(template)
+
+    map_name = "TestMap"
+    first_key = "FirstKey"
+    second_key = "SecondKey"
+    expected = "ExpectedValue"
+
+    values = [map_name, first_key, second_key, {"DefaultValue": "Default"}]
+
+    template.template["Mappings"] = {map_name: {first_key: {second_key: expected}}}
+
+    result = functions.enhanced_find_in_map(template, values)
+
+    assert result == expected
+
+    values = [map_name, first_key, "FakeKey", {"DefaultValue": "Default"}]
+
+    result = functions.enhanced_find_in_map(template, values)
+
+    assert result == "Default"
+
+
 def test_get_att():
     template = {"Resources": {}}