Adding TransferFamilyAuthorizer and TransferFamilyAuthorizerResponse class

leandrodamascena · leandrodamascena · commit 0681ad90461e · 2025-01-23T22:23:20.000Z
diff --git a/aws_lambda_powertools/utilities/data_classes/__init__.py b/aws_lambda_powertools/utilities/data_classes/__init__.py
@@ -45,6 +45,7 @@
 from .ses_event import SESEvent
 from .sns_event import SNSEvent
 from .sqs_event import SQSEvent
+from .transfer_family_event import TransferFamilyAuthorizer, TransferFamilyAuthorizerResponse
 from .vpc_lattice import VPCLatticeEvent, VPCLatticeEventV2
 
 __all__ = [
@@ -87,4 +88,6 @@
     "VPCLatticeEvent",
     "VPCLatticeEventV2",
     "CloudFormationCustomResourceEvent",
+    "TransferFamilyAuthorizerResponse",
+    "TransferFamilyAuthorizer",
 ]
diff --git a/aws_lambda_powertools/utilities/data_classes/transfer_family_event.py b/aws_lambda_powertools/utilities/data_classes/transfer_family_event.py
@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from aws_lambda_powertools.utilities.data_classes.common import (
+    DictWrapper,
+)
+
+
+class TransferFamilyAuthorizer(DictWrapper):
+    @property
+    def username(self) -> str:
+        """The username used for authentication"""
+        return self["username"]
+
+    @property
+    def password(self) -> str | None:
+        """
+        The password used for authentication.
+        None in case customer authenticating with certificates
+        """
+        return self["password"]
+
+    @property
+    def protocol(self) -> str:
+        """The protocol can be SFTP, FTP or FTPS"""
+        return self["protocol"]
+
+    @property
+    def server_id(self) -> str:
+        """The AWS Transfer Family ServerID"""
+        return self["serverId"]
+
+    @property
+    def source_ip(self) -> str:
+        """The customer IP used for connection"""
+        return self["sourceIp"]
+
+
+class TransferFamilyAuthorizerResponse:
+
+    def _build_authentication_response(
+        self,
+        role_arn: str,
+        policy: str | None = None,
+        home_directory: str | None = None,
+        home_directory_details: dict | None = None,
+        home_directory_type: Literal["LOGICAL", "PATH"] = "PATH",
+        user_gid: int | None = None,
+        user_uid: int | None = None,
+        public_keys: str | None = None,
+    ) -> dict[str, Any]:
+
+        response: dict[str, Any] = {}
+
+        if home_directory_type == "PATH":
+            if not home_directory:
+                raise ValueError("home_directory must be set when home_directory_type is PATH")
+
+            response["HomeDirectory"] = home_directory
+        elif home_directory_type == "LOGICAL":
+            if not home_directory_details:
+                raise ValueError("home_directory_details must be set when home_directory_type is LOGICAL")
+
+            response["HomeDirectoryDetails"] = [home_directory_details]
+
+        else:
+            raise ValueError(f"Invalid home_directory_type: {home_directory_type}")
+
+        if user_uid is not None:
+            response["PosixProfile"] = {"Gid": user_gid, "Uid": user_gid}
+
+        if policy:
+            response["Policy"] = policy
+
+        if public_keys:
+            response["PublicKeys"] = public_keys
+
+        response["Role"] = role_arn
+        response["HomeDirectoryType"] = home_directory_type
+
+        return response
+
+    def build_authentication_response_efs(
+        self,
+        role_arn: str,
+        user_gid: int,
+        user_uid: int,
+        policy: str | None = None,
+        home_directory: str | None = None,
+        home_directory_details: dict | None = None,
+        home_directory_type: Literal["LOGICAL", "PATH"] = "PATH",
+        public_keys: str | None = None,
+    ) -> dict[str, Any]:
+        """
+        Build an authentication response for AWS Transfer Family using EFS (Elastic File System).
+
+        Parameters:
+        -----------
+        role_arn : str
+            The Amazon Resource Name (ARN) of the IAM role.
+        user_gid : int
+            The group ID of the user.
+        user_uid : int
+            The user ID.
+        policy : str | None, optional
+            The IAM policy document. Defaults to None.
+        home_directory : str | None, optional
+            The home directory path. Required if home_directory_type is "PATH". Defaults to None.
+        home_directory_details : dict | None, optional
+            Details of the home directory. Required if home_directory_type is "LOGICAL". Defaults to None.
+        home_directory_type : Literal["LOGICAL", "PATH"], optional
+            The type of home directory. Must be either "LOGICAL" or "PATH". Defaults to "PATH".
+        public_keys : str | None, optional
+            The public keys associated with the user. Defaults to None.
+
+        Returns:
+        --------
+        dict[str, Any]
+            A dictionary containing the authentication response with various details such as
+            role ARN, policy, home directory information, and user details.
+
+        Raises:
+        -------
+        ValueError
+            If an invalid home_directory_type is provided or if required parameters are missing
+            for the specified home_directory_type.
+        """
+
+        return self._build_authentication_response(
+            role_arn=role_arn,
+            policy=policy,
+            home_directory=home_directory,
+            home_directory_details=home_directory_details,
+            home_directory_type=home_directory_type,
+            public_keys=public_keys,
+            user_gid=user_gid,
+            user_uid=user_uid,
+        )
+
+    def build_authentication_response_s3(
+        self,
+        role_arn: str,
+        policy: str | None = None,
+        home_directory: str | None = None,
+        home_directory_details: dict | None = None,
+        home_directory_type: Literal["LOGICAL", "PATH"] = "PATH",
+        public_keys: str | None = None,
+    ) -> dict[str, Any]:
+        """
+        Build an authentication response for Amazon S3.
+
+        This method constructs an authentication response tailored for S3 access,
+        likely by calling an internal method with the provided parameters.
+
+        Parameters:
+        -----------
+        role_arn : str
+            The Amazon Resource Name (ARN) of the IAM role for S3 access.
+        policy : str | None, optional
+            The IAM policy document for S3 access. Defaults to None.
+        home_directory : str | None, optional
+            The home directory path in S3. Required if home_directory_type is "PATH". Defaults to None.
+        home_directory_details : dict | None, optional
+            Details of the home directory in S3. Required if home_directory_type is "LOGICAL". Defaults to None.
+        home_directory_type : Literal["LOGICAL", "PATH"], optional
+            The type of home directory in S3. Must be either "LOGICAL" or "PATH". Defaults to "PATH".
+        public_keys : str | None, optional
+            The public keys associated with the user for S3 access. Defaults to None.
+
+        Returns:
+        --------
+        dict[str, Any]
+            A dictionary containing the authentication response with various details such as
+            role ARN, policy, home directory information, and potentially other S3-specific attributes.
+
+        Raises:
+        -------
+        ValueError
+            If an invalid home_directory_type is provided or if required parameters are missing
+            for the specified home_directory_type.
+        """
+        return self._build_authentication_response(
+            role_arn=role_arn,
+            policy=policy,
+            home_directory=home_directory,
+            home_directory_details=home_directory_details,
+            home_directory_type=home_directory_type,
+            public_keys=public_keys,
+        )
diff --git a/aws_lambda_powertools/utilities/parser/models/__init__.py b/aws_lambda_powertools/utilities/parser/models/__init__.py
@@ -109,7 +109,7 @@
 )
 from .sns import SnsModel, SnsNotificationModel, SnsRecordModel
 from .sqs import SqsAttributesModel, SqsModel, SqsMsgAttributeModel, SqsRecordModel
-from .transfer_family import TransferFamily
+from .transfer_family import TransferFamilyAuthorizer
 from .vpc_lattice import VpcLatticeModel
 from .vpc_latticev2 import VpcLatticeV2Model
 
@@ -180,7 +180,7 @@
     "SqsAttributesModel",
     "S3SqsEventNotificationModel",
     "S3SqsEventNotificationRecordModel",
-    "TransferFamily",
+    "TransferFamilyAuthorizer",
     "APIGatewayProxyEventModel",
     "APIGatewayEventRequestContext",
     "APIGatewayEventAuthorizer",
diff --git a/aws_lambda_powertools/utilities/parser/models/transfer_family.py b/aws_lambda_powertools/utilities/parser/models/transfer_family.py
@@ -4,7 +4,7 @@
 from pydantic.networks import IPvAnyAddress
 
 
-class TransferFamily(BaseModel):
+class TransferFamilyAuthorizer(BaseModel):
     username: str
     password: Optional[str] = None
     protocol: Literal["SFTP", "FTP", "FTPS"]
diff --git a/docs/utilities/data_classes.md b/docs/utilities/data_classes.md
@@ -108,6 +108,8 @@ Log Data Event for Troubleshooting
 | [SES](#ses)                                                                   | `SESEvent`                                         |
 | [SNS](#sns)                                                                   | `SNSEvent`                                         |
 | [SQS](#sqs)                                                                   | `SQSEvent`                                         |
+| [TransferFamilyAuthorizer]                                                    | `TransferFamilyAuthorizer`                         |
+| [TransferFamilyAuthorizerResponse]                                            | `TransferFamilyAuthorizerResponse`                 |
 | [VPC Lattice V2](#vpc-lattice-v2)                                             | `VPCLatticeV2Event`                                |
 | [VPC Lattice V1](#vpc-lattice-v1)                                             | `VPCLatticeEvent`                                  |
 
diff --git a/docs/utilities/parser.md b/docs/utilities/parser.md
@@ -132,7 +132,7 @@ The example above uses `SqsModel`. Other built-in models can be found below.
 | **SesModel**                                | Lambda Event Source payload for Amazon Simple Email Service                           |
 | **SnsModel**                                | Lambda Event Source payload for Amazon Simple Notification Service                    |
 | **SqsModel**                                | Lambda Event Source payload for Amazon SQS                                            |
-| **TransferFamily**                          | Lambda Event Source payload for AWS Transfer Family custom identity provider          |
+| **TransferFamilyAuthorizer**                | Lambda Event Source payload for AWS Transfer Family Lambda authorizer                 |
 | **VpcLatticeModel**                         | Lambda Event Source payload for Amazon VPC Lattice                                    |
 | **VpcLatticeV2Model**                       | Lambda Event Source payload for Amazon VPC Lattice v2 payload                         |
 
diff --git a/tests/events/TransferFamilyAuthorizer.json b/tests/events/TransferFamilyAuthorizer.json
diff --git a/tests/unit/data_classes/required_dependencies/test_transfer_family_event.py b/tests/unit/data_classes/required_dependencies/test_transfer_family_event.py
@@ -0,0 +1,142 @@
+import pytest
+
+from aws_lambda_powertools.utilities.data_classes.transfer_family_event import (
+    TransferFamilyAuthorizer,
+    TransferFamilyAuthorizerResponse,
+)
+from tests.functional.utils import load_event
+
+
+def test_transfer_family_event():
+    raw_event = load_event("transferFamilyAuthorizer.json")
+    parsed_event = TransferFamilyAuthorizer(raw_event)
+
+    assert parsed_event.username == raw_event["username"]
+    assert parsed_event.password == raw_event["password"]
+    assert parsed_event.protocol == raw_event["protocol"]
+    assert parsed_event.server_id == raw_event["serverId"]
+    assert parsed_event.source_ip == raw_event["sourceIp"]
+
+
+@pytest.mark.parametrize("home_directory_type", ["LOGICAL", "PATH"])
+def test_build_authentication_response_s3(home_directory_type):
+
+    # GIVEN a Authorizer response
+    response = TransferFamilyAuthorizerResponse()
+
+    role_arn = "arn:aws:iam::123456789012:role/S3Access"
+    policy = '{"Version": "2012-10-17", "Statement": [{"Effect": "Allow", "Action": "s3:*", "Resource": "*"}]}'
+    home_directory = "/bucket/user" if home_directory_type == "PATH" else None
+    home_directory_details = (
+        {"Entry": "/", "Target": "/bucket/${transfer:UserName}"} if home_directory_type == "LOGICAL" else None
+    )
+    public_keys = "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC0g+Z"
+
+    # WHEN building an authentication response for S3 with different home directory types
+    response = response.build_authentication_response_s3(
+        role_arn=role_arn,
+        policy=policy,
+        home_directory=home_directory,
+        home_directory_details=home_directory_details,
+        home_directory_type=home_directory_type,
+        public_keys=public_keys,
+    )
+
+    # THEN the authentication response is correctly built
+    assert isinstance(response, dict)
+    assert response.get("Role") == role_arn
+    assert response.get("Policy") == policy
+    assert response.get("PublicKeys") == public_keys
+
+    if home_directory_type == "PATH":
+        assert response.get("HomeDirectory") == home_directory
+        assert "HomeDirectoryDetails" not in response
+    else:
+        assert response.get("HomeDirectoryDetails") == [home_directory_details]
+        assert "HomeDirectory" not in response
+
+
+@pytest.mark.parametrize("home_directory_type", ["LOGICAL", "PATH"])
+def test_build_authentication_response_efs(home_directory_type):
+
+    # GIVEN a Authorizer response
+    response = TransferFamilyAuthorizerResponse()
+
+    role_arn = "arn:aws:iam::123456789012:role/S3Access"
+    home_directory = "/bucket/user" if home_directory_type == "PATH" else None
+    home_directory_details = (
+        {"Entry": "/", "Target": "/bucket/${transfer:UserName}"} if home_directory_type == "LOGICAL" else None
+    )
+
+    # WHEN building an authentication response for EFS with different home directory types
+    response = response.build_authentication_response_efs(
+        role_arn=role_arn,
+        home_directory=home_directory,
+        home_directory_details=home_directory_details,
+        home_directory_type=home_directory_type,
+        user_gid=0,
+        user_uid=0,
+    )
+
+    # THEN the authentication response is correctly built
+    assert isinstance(response, dict)
+    assert response.get("Role") == role_arn
+
+    if home_directory_type == "PATH":
+        assert response.get("HomeDirectory") == home_directory
+        assert "HomeDirectoryDetails" not in response
+    else:
+        assert response.get("HomeDirectoryDetails") == [home_directory_details]
+        assert "HomeDirectory" not in response
+
+
+def test_build_authentication_missing_home_directory():
+
+    # GIVEN a Authorizer response
+    response = TransferFamilyAuthorizerResponse()
+
+    # WHEN home_directory_details is empty and type is LOGICAL
+    role_arn = "arn:aws:iam::123456789012:role/S3Access"
+    home_directory_details = {}
+    home_directory_type = "LOGICAL"
+
+    # THEN must raise an exception
+    with pytest.raises(ValueError):
+        response = response.build_authentication_response_efs(
+            role_arn=role_arn,
+            home_directory_details=home_directory_details,
+            home_directory_type=home_directory_type,
+            user_gid=0,
+            user_uid=0,
+        )
+
+
+def test_build_authentication_response_invalid_type():
+    # GIVEN a Authorizer response
+    response = TransferFamilyAuthorizerResponse()
+
+    # WHEN set an invalid home_directory_type
+    invalid_type = "INVALID"
+
+    # THEN must raise an exception
+    with pytest.raises(ValueError):
+        response.build_authentication_response_s3(
+            role_arn="arn:aws:iam::123456789012:role/S3Access",
+            home_directory_type=invalid_type,
+        )
+
+
+def test_build_authentication_response_missing_required_params():
+    # GIVEN a Authorizer response
+    response = TransferFamilyAuthorizerResponse()
+
+    # WHEN set a PATH without home_directory
+    home_directory_type = "PATH"
+
+    # THEN must raise an exception
+    with pytest.raises(ValueError):
+        response.build_authentication_response_s3(
+            role_arn="arn:aws:iam::123456789012:role/S3Access",
+            home_directory_type=home_directory_type,
+            # Missing required home_directory for PATH type
+        )
diff --git a/tests/unit/parser/_pydantic/test_transfer_family.py b/tests/unit/parser/_pydantic/test_transfer_family.py
