diff --git a/templates/api/handler.py b/templates/api/handler.py index 264793c..cbe4ee6 100644 --- a/templates/api/handler.py +++ b/templates/api/handler.py @@ -1,5 +1,3 @@ -from json import dumps - from aws_lambda_powertools import Logger, Metrics, Tracer from aws_lambda_powertools.event_handler import APIGatewayRestResolver from aws_lambda_powertools.event_handler.api_gateway import Response @@ -7,6 +5,7 @@ from pydantic import ValidationError from templates.api.models import Item +from templates.api.response import JsonResponse from templates.api.settings import Settings from templates.repository import Repository @@ -32,17 +31,14 @@ def get_item(id: str) -> Response: """ try: if (item := repository.get_item(id)) is None: - return Response(status_code=404, content_type="application/json", body=dumps({"message": "Not found"})) + return JsonResponse({"message": "Not found"}, status_code=404) item = Item.model_validate(item) # Validate model after retrieval to ensure data integrity except Exception as exc: - message = ( - "Item validation failed" if isinstance(exc, ValidationError) else f"Error retrieving item with id {id}" - ) - logger.error(message, exc_info=exc) - return Response( - status_code=500, content_type="application/json", body=dumps({"message": "Internal server error"}) - ) - return Response(status_code=200, content_type="application/json", body=item.dump()) + message = "Item validation failed" if isinstance(exc, ValidationError) else "Error retrieving item" + logger.error(message, exc_info=exc, extra={"itemId": id}) + return JsonResponse({"message": "Internal server error"}, status_code=500) + + return JsonResponse(item.dump()) @app.post("/items") @@ -55,17 +51,15 @@ def create_item() -> Response: try: item = Item.model_validate_json(app.current_event.body) except ValidationError as exc: - return Response(status_code=422, content_type="application/json", body=dumps({"errors": exc.errors()})) + return JsonResponse({"errors": exc.errors()}, status_code=422) try: repository.put_item(item.model_dump()) except Exception as exc: - logger.error("DynamoDB put_item failed", exc_info=exc) - return Response( - status_code=500, content_type="application/json", body=dumps({"message": "Internal server error"}) - ) + logger.error("DynamoDB put_item failed", exc_info=exc, extra={"itemId": item.id}) + return JsonResponse({"message": "Internal server error"}, status_code=500) - return Response(status_code=201, content_type="application/json", body=item.dump()) + return JsonResponse(item.dump(), status_code=201) @logger.inject_lambda_context diff --git a/templates/api/response.py b/templates/api/response.py new file mode 100644 index 0000000..161ee56 --- /dev/null +++ b/templates/api/response.py @@ -0,0 +1,27 @@ +from json import dumps + +from aws_lambda_powertools.event_handler.api_gateway import Response + +SECURITY_HEADERS = { + "X-Content-Type-Options": "nosniff", + "X-Frame-Options": "DENY", + "Strict-Transport-Security": "max-age=31536000; includeSubDomains", +} + + +class JsonResponse(Response): + """An HTTP response with JSON body and security headers.""" + + def __init__(self, body: dict | str, status_code: int = 200) -> None: + """Initialize the JSON response. + + Args: + body: The response body as a dictionary or JSON string. + status_code: The HTTP status code. + """ + super().__init__( + status_code=status_code, + body=body if isinstance(body, str) else dumps(body), + content_type="application/json", + headers=SECURITY_HEADERS, + ) diff --git a/tests/api/test_handler.py b/tests/api/test_handler.py index 225fe41..19c3a78 100644 --- a/tests/api/test_handler.py +++ b/tests/api/test_handler.py @@ -192,5 +192,21 @@ def test_get_item_validation_error_returns_500(mock_repo, lambda_context): assert body["message"] == "Internal server error" +def test_response_headers_contain_security_headers(mock_repo, lambda_context): + """API responses should contain standard security headers.""" + from templates.api.handler import main + + mock_repo.get_item.return_value = {"id": "abc", "name": "Widget"} + + event = _apigw_event("GET", "/items/abc", path_params={"id": "abc"}) + response = main(event, lambda_context) + + assert response["statusCode"] == 200 + headers = response["multiValueHeaders"] + assert headers["X-Content-Type-Options"] == ["nosniff"] + assert headers["X-Frame-Options"] == ["DENY"] + assert "Strict-Transport-Security" in headers + + if __name__ == "__main__": main()