Repository: getmoto/moto
Pull request #7942 (2024-08-07 10:03:20)

Core: Allow config of whitelisted services

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/getmoto/moto/pull/7942?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 94.35%. Comparing base [(`2a80170`)](https://app.codecov.io/gh/getmoto/moto/commit/2a8017076e0a4669f2e2e55e2198aadacb13bb83?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) to head [(`b317a91`)](https://app.codecov.io/gh/getmoto/moto/commit/b317a915df8b63f138ed8a428043cb354efefdbd?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #7942   +/-   ##
=======================================
  Coverage   94.35%   94.35%           
=======================================
  Files        1120     1120           
  Lines       95956    95968   +12     
=======================================
+ Hits        90542    90554   +12     
  Misses       5414     5414           
```

| [Flag](https://app.codecov.io/gh/getmoto/moto/pull/7942/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | Coverage Δ | |
|---|---|---|
| [servertests](https://app.codecov.io/gh/getmoto/moto/pull/7942/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | `28.96% <50.00%> (+<0.01%)` | :arrow_up: |
| [unittests](https://app.codecov.io/gh/getmoto/moto/pull/7942/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | `94.32% <100.00%> (+<0.01%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto#carryforward-flags-in-the-pull-request-comment) to find out more.

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/getmoto/moto/pull/7942?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).


Patch:
diff --git a/docs/docs/configuration/index.rst b/docs/docs/configuration/index.rst
index 9466d7497..ba0b50358 100644
--- a/docs/docs/configuration/index.rst
+++ b/docs/docs/configuration/index.rst
@@ -20,6 +20,7 @@ If you are using the decorators, some options are configurable within the decora
                 "services": ["dynamodb"]
             },
             "reset_boto3_session": True,
+            "service_whitelist": None,
         },
         "iam": {"load_aws_managed_policies": False},
         "stepfunctions": {"execute_state_machine": True},
@@ -52,6 +53,17 @@ That is why Moto resets the `boto3-Session`, to make sure that it is recreated w
 
 If all of your tests use Moto, and you never want to reach out to AWS, you can choose to _not_ reset the `boto3-session`. New boto3-clients that are created will reuse the `boto3-Session` (with fake credentials), making Moto much more performant.
 
+Whitelist Services
+-------------------
+The `mock_aws` decorator will allow requests to all supported services. If you want to restrict this to only specific services, you can configure a service whitelist.
+For example, if you only want your application to use DynamoDB and S3:
+
+```
+mock_aws(config={"core": {"service_whitelist": ["dynamodb", "s3"]}})
+```
+
+If the application under test tries to access any other services, Moto will throw a `ServiceNotWhitelisted`-exception.
+
 AWS Managed Policies
 --------------------
 Moto comes bundled with all Managed Policies that AWS exposes, which are updated regularly. However, they are not loaded unless specifically requested for performance reasons.
diff --git a/moto/awslambda/models.py b/moto/awslambda/models.py
index ce66fef66..7943659cc 100644
--- a/moto/awslambda/models.py
+++ b/moto/awslambda/models.py
@@ -1150,9 +1150,9 @@ class LambdaFunction(CloudFormationModel, DockerModel):
         spec = {
             "Code": properties["Code"],
             "FunctionName": resource_name,
-            "Handler": properties["Handler"],
+            "Handler": properties.get("Handler"),
             "Role": properties["Role"],
-            "Runtime": properties["Runtime"],
+            "Runtime": properties.get("Runtime"),
         }
 
         # NOTE: Not doing `properties.get(k, DEFAULT)` to avoid duplicating the
diff --git a/moto/core/botocore_stubber.py b/moto/core/botocore_stubber.py
index 929326168..0872fc8f3 100644
--- a/moto/core/botocore_stubber.py
+++ b/moto/core/botocore_stubber.py
@@ -7,7 +7,8 @@ from botocore.awsrequest import AWSResponse
 import moto.backend_index as backend_index
 from moto.core.base_backend import BackendDict
 from moto.core.common_types import TYPE_RESPONSE
-from moto.core.config import passthrough_service, passthrough_url
+from moto.core.config import passthrough_service, passthrough_url, service_whitelisted
+from moto.core.exceptions import ServiceNotWhitelisted
 from moto.core.utils import get_equivalent_url_in_aws_domain
 
 
@@ -56,6 +57,9 @@ class BotocoreStubber:
                 if passthrough_service(service):
                     return None
 
+                if not service_whitelisted(service):
+                    raise ServiceNotWhitelisted(service)
+
                 import moto.backends as backends
                 from moto.core import DEFAULT_ACCOUNT_ID
                 from moto.core.exceptions import HTTPException
diff --git a/moto/core/config.py b/moto/core/config.py
index 1b6cfa72b..47a3e27ee 100644
--- a/moto/core/config.py
+++ b/moto/core/config.py
@@ -1,5 +1,5 @@
 import re
-from typing import List, TypedDict
+from typing import List, Optional, TypedDict
 
 
 class _docker_config(TypedDict, total=False):
@@ -15,6 +15,7 @@ class _core_config(TypedDict, total=False):
     mock_credentials: bool
     passthrough: _passthrough_config
     reset_boto3_session: bool
+    service_whitelist: Optional[List[str]]
 
 
 class _iam_config(TypedDict, total=False):
@@ -49,6 +50,7 @@ default_user_config: DefaultConfig = {
         "mock_credentials": True,
         "passthrough": {"urls": [], "services": []},
         "reset_boto3_session": True,
+        "service_whitelist": None,
     },
     "iam": {"load_aws_managed_policies": False},
     "stepfunctions": {"execute_state_machine": False},
@@ -56,6 +58,11 @@ default_user_config: DefaultConfig = {
 }
 
 
+def service_whitelisted(service: str) -> bool:
+    services_whitelisted = default_user_config.get("core", {}).get("service_whitelist")
+    return services_whitelisted is None or service in services_whitelisted
+
+
 def passthrough_service(service: str) -> bool:
     passthrough_services = (
         default_user_config.get("core", {}).get("passthrough", {}).get("services", [])
diff --git a/moto/core/exceptions.py b/moto/core/exceptions.py
index 69a5ae457..86595bc8e 100644
--- a/moto/core/exceptions.py
+++ b/moto/core/exceptions.py
@@ -200,3 +200,15 @@ class InvalidToken(AWSError):
 
     def __init__(self, message: str = "Invalid token"):
         super().__init__(f"Invalid Token: {message}", "InvalidToken")
+
+
+class ServiceNotWhitelisted(Exception):
+    def __init__(self, service_name: str):
+        from moto.settings import default_user_config
+
+        services_whitelisted = default_user_config.get("core", {}).get(
+            "service_whitelist"
+        )
+        super().__init__(
+            f"Service {service_name} not whitelisted. Only services {services_whitelisted} are allowed."
+        )


Test patch:
diff --git a/tests/test_awslambda/test_awslambda_cloudformation.py b/tests/test_awslambda/test_awslambda_cloudformation.py
index 09e3edf8c..6e1198a0f 100644
--- a/tests/test_awslambda/test_awslambda_cloudformation.py
+++ b/tests/test_awslambda/test_awslambda_cloudformation.py
@@ -1,4 +1,5 @@
 import io
+import json
 import re
 import zipfile
 from string import Template
@@ -54,6 +55,26 @@ template = Template(
 }"""
 )
 
+
+code_image_template = {
+    "AWSTemplateFormatVersion": "2010-09-09",
+    "Resources": {
+        "LambdaFunction": {
+            "Type": "AWS::Lambda::Function",
+            "Properties": {
+                "FunctionName": "ContainerLambdaFunction",
+                "Role": "dummy-role-arn",
+                "Code": {
+                    "ImageUri": "123456789012.dkr.ecr.us-west-2.amazonaws.com/my-lambda-container:latest"
+                },
+                "PackageType": "Image",
+                "MemorySize": 128,
+                "Timeout": 30,
+            },
+        }
+    },
+}
+
 event_source_mapping_template = Template(
     """{
     "AWSTemplateFormatVersion": "2010-09-09",
@@ -114,6 +135,25 @@ def test_lambda_can_be_deleted_by_cloudformation():
     assert e.value.response["Error"]["Code"] == "ResourceNotFoundException"
 
 
+@mock_aws
+def test_lambda_with_code_image_can_be_created():
+    our_template = code_image_template.copy()
+    our_template["Resources"]["LambdaFunction"]["Properties"]["Role"] = get_role_arn()
+    stack_name = f"stack{str(uuid4())[0:6]}"
+    cf = boto3.client("cloudformation", region_name="us-east-1")
+    cf.create_stack(StackName=stack_name, TemplateBody=json.dumps(code_image_template))
+
+    lmbda = boto3.client("lambda", region_name="us-east-1")
+    functions = lmbda.list_functions()["Functions"]
+    function_names = [f["FunctionName"] for f in functions]
+    assert "ContainerLambdaFunction" in function_names
+
+    cf.delete_stack(StackName=stack_name)
+    functions = lmbda.list_functions()["Functions"]
+    function_names = [f["FunctionName"] for f in functions]
+    assert "ContainerLambdaFunction" not in function_names
+
+
 @mock_aws
 def test_event_source_mapping_create_from_cloudformation_json():
     sqs = boto3.resource("sqs", region_name="us-east-1")
diff --git a/tests/test_core/test_config.py b/tests/test_core/test_config.py
index 63ba387ff..20b2d5939 100644
--- a/tests/test_core/test_config.py
+++ b/tests/test_core/test_config.py
@@ -1,3 +1,7 @@
+from unittest import SkipTest
+
+import boto3
+import pytest
 import requests
 
 from moto import mock_aws, settings
@@ -6,6 +10,7 @@ from moto.awslambda.utils import get_backend
 from moto.awslambda_simple.models import LambdaSimpleBackend
 from moto.core import DEFAULT_ACCOUNT_ID
 from moto.core.config import default_user_config
+from moto.core.exceptions import ServiceNotWhitelisted
 
 
 @mock_aws
@@ -50,3 +55,33 @@ def test_change_configuration_using_api() -> None:
         f"http://{base_url}/moto-api/config",
         json={"batch": {"use_docker": True}, "lambda": {"use_docker": True}},
     )
+
+
+def test_whitelist() -> None:
+    if not settings.TEST_DECORATOR_MODE:
+        raise SkipTest("No point in testing this anywhere else")
+
+    with mock_aws(config={"core": {"service_whitelist": ["s3"]}}):
+        dynamodb = boto3.client("dynamodb", "us-east-1")
+        s3 = boto3.client("s3", "us-east-1")
+
+        # S3 is whitelisted, so all calls are allowed
+        s3.list_buckets()
+
+        # DynamoDB is not whitelisted
+        with pytest.raises(ServiceNotWhitelisted):
+            dynamodb.list_tables()
+
+    # Empty whitelist, so no calls are allowed
+    with mock_aws(config={"core": {"service_whitelist": []}}):
+        with pytest.raises(ServiceNotWhitelisted):
+            s3.list_buckets()
+
+        with pytest.raises(ServiceNotWhitelisted):
+            dynamodb.list_tables()
+
+    # Both are whitelisted, so both are allowed
+    with mock_aws(config={"core": {"service_whitelist": ["dynamodb", "s3"]}}):
+        s3.list_buckets()
+
+        dynamodb.list_tables()