🛡️ Sentinel: Enhance input validation and error sanitization

templates/agent/handler.py

@@ -2,9 +2,11 @@
 from aws_lambda_powertools.event_handler import BedrockAgentFunctionResolver
 from aws_lambda_powertools.utilities.data_classes import BedrockAgentEvent
 from aws_lambda_powertools.utilities.typing import LambdaContext
+from pydantic import ValidationError
 
 from templates.agent.models import Item
 from templates.agent.settings import Settings
+from templates.models import Entity
 from templates.repository import Repository
 
 settings = Settings()  # type: ignore
@@ -29,11 +31,20 @@ def get_item(item_id: str) -> dict:
         The item details or an error message.
     """
     logger.info("Retrieving item", extra={"itemId": item_id})
+    try:
+        Entity(id=item_id)
+    except ValidationError:
+        logger.warning("Invalid item ID provided", extra={"itemId": item_id})
+        return {"error": f"Invalid item ID '{item_id}'"}
+
     try:
         item = repository.get_item(item_id)
         if not item:
             return {"error": f"Item {item_id} not found"}
         return Item.model_validate(item).dump()
+    except ValidationError as error:
+        logger.error("Item validation failed", extra={"itemId": item_id}, exc_info=error)
+        return {"error": "Internal server error"}
     except Exception as error:
         logger.error("Failed to get item", extra={"itemId": item_id}, exc_info=error)
         return {"error": f"Failed to get item with ID '{item_id}'"}
@@ -57,6 +68,9 @@ def create_item(item_id: str, name: str, description: str | None = None) -> dict
         item = Item(id=item_id, name=name, description=description).dump()
         repository.put_item(item)
         return item
+    except ValidationError as error:
+        logger.warning("Invalid item data provided", extra={"itemId": item_id}, exc_info=error)
+        return {"error": "Invalid item data"}
     except Exception as error:
         logger.error("Failed to create item", extra={"itemId": item_id}, exc_info=error)
         return {"error": f"Failed to create item with ID '{item_id}'"}

templates/api/handler.py

@@ -6,8 +6,8 @@
 
 from templates.api.models import Item
 from templates.api.response import JsonResponse
-from templates.models import Entity
 from templates.api.settings import Settings
+from templates.models import Entity
 from templates.repository import Repository
 
 settings = Settings()
@@ -33,7 +33,7 @@ def get_item(id: str) -> Response:
     try:
         Entity(id=id)
     except ValidationError:
-        return JsonResponse({"message": "Invalid item ID length"}, status_code=400)
+        return JsonResponse({"message": "Invalid item ID"}, status_code=400)
 
     try:
         if (item := repository.get_item(id)) is None:

templates/graphql/handler.py

@@ -6,6 +6,7 @@
 
 from templates.graphql.models import Item
 from templates.graphql.settings import Settings
+from templates.models import Entity
 from templates.repository import Repository
 
 settings = Settings()  # type: ignore
@@ -29,13 +30,21 @@ def get_item(id: str) -> dict | None:
     Returns:
         The item if found, or None.
     """
+    try:
+        Entity(id=id)
+    except ValidationError:
+        logger.warning("Invalid item ID provided", extra={"itemId": id})
+        raise RuntimeError(f"Invalid item ID '{id}'") from None
+
     try:
         if (item := repository.get_item(id)) is None:
             return None
         return Item.model_validate(item).dump()
     except Exception as error:
-        logger.error("Failed to get item", extra={"itemId": id}, exc_info=error)
-        raise RuntimeError(f"Failed to get item with ID '{id}'") from None
+        is_val_error = isinstance(error, ValidationError)
+        message = "Item validation failed" if is_val_error else f"Failed to get item with ID '{id}'"
+        logger.error(message, extra={"itemId": id}, exc_info=error)
+        raise RuntimeError(message) from None
 
 
 @app.resolver(type_name="Query", field_name="listItems")
@@ -68,7 +77,10 @@ def create_item(name: str) -> dict:
         item = Item(name=name).dump()
         repository.put_item(item)
         return item
-    except (ValidationError, Exception) as error:
+    except ValidationError as error:
+        logger.warning("Invalid item data provided", extra={"itemName": name}, exc_info=error)
+        raise RuntimeError("Invalid item data") from None
+    except Exception as error:
         logger.error("Failed to create item", extra={"itemName": name}, exc_info=error)
         raise RuntimeError(f"Failed to create item with name '{name}'") from None
 

templates/models.py

@@ -29,4 +29,5 @@ class Entity(Object):
         default_factory=lambda: str(uuid4()),
         min_length=1,
         max_length=50,
+        pattern=r"^[a-zA-Z0-9-_]+$",
     )

templates/sqs/handler.py

@@ -43,7 +43,7 @@ def handle_record(self, record: SQSRecord) -> None:
         try:
             message = SqsMessage.model_validate_json(record.body)
             processed = ProcessedItem(id=message.id, content=message.content, status="PROCESSED")
-            self._repository.put_item(processed.model_dump())
+            self._repository.put_item(processed.dump())
             logger.info("Successfully processed and stored message", extra={"messageId": message.id})
         except Exception as exc:
             logger.error("Failed to process SQS record", exc_info=exc)

tests/agent/test_handler.py

@@ -47,6 +47,15 @@ def test_handler_get_item_not_found():
     assert "not found" in result["error"]
 
 
+def test_handler_get_item_invalid_id():
+    from templates.agent.handler import get_item
+
+    result = get_item("invalid!")
+
+    assert "error" in result
+    assert "Invalid item ID" in result["error"]
+
+
 def test_handler_create_item(repository):
     from templates.agent.handler import create_item
 

tests/api/test_handler.py

@@ -109,7 +109,19 @@ def test_get_item_id_too_long(mock_repo, lambda_context):
 
     assert response["statusCode"] == 400
     body = loads(response["body"])
-    assert body["message"] == "Invalid item ID length"
+    assert body["message"] == "Invalid item ID"
+
+
+def test_get_item_id_invalid_pattern(mock_repo, lambda_context):
+    """GET /items/{id} returns 400 when the ID contains invalid characters."""
+    import templates.api.handler as handler_module
+
+    event = _apigw_event("GET", "/items/invalid!", path_params={"id": "invalid!"})
+    response = handler_module.main(event, lambda_context)
+
+    assert response["statusCode"] == 400
+    body = loads(response["body"])
+    assert body["message"] == "Invalid item ID"
 
 
 def test_post_item_invalid_body(mock_repo, lambda_context):

tests/graphql/test_handler.py

@@ -58,6 +58,17 @@ def test_sensitive_data_exposure(repository, lambda_context):
     assert "internal_secret" not in result_list[0]
 
 
+def test_get_item_invalid_id(lambda_context):
+    import pytest
+
+    from templates.graphql.handler import main
+
+    event = {"info": {"parentTypeName": "Query", "fieldName": "getItem"}, "arguments": {"id": "invalid!"}}
+    with pytest.raises(RuntimeError) as excinfo:
+        main(event, lambda_context)
+    assert "Invalid item ID" in str(excinfo.value)
+
+
 def test_error_message_information_leakage(lambda_context, mocker):
     """Verify that internal error details are NOT leaked to the client."""
     from templates.graphql import handler