Add support for ec2.assume_role command (#43)

* Add support for ec2.assume_role command

* update changelog

* lint

* update changelog

Author: blink1073

CHANGELOG.md

@@ -1,5 +1,8 @@
 # Changelog
 
+## 3.7.0 - 2025-03-18
+- Add support for `ec2.assume_role` to the list of `AvailableCommands`.
+
 ## 3.6.0 - 2024-11-18
 - Address Python 3.12+ compatibility issues.
 - Address pydantic v2 compatibility issues.

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "shrub.py"
-version = "3.6.0"
+version = "3.7.0"
 description = "Library for creating evergreen configurations"
 authors = ["DevProd Services & Integrations Team <[email protected]>"]
 license = "Apache-2.0"

src/shrub/v3/evg_command.py

@@ -14,6 +14,7 @@
     "attach.results",
     "attach.xunit_results",
     "downstream_expansions.set",
+    "ec2.assume_role",
     "expansions.update",
     "expansions.write",
     "generate.tasks",
@@ -350,6 +351,31 @@ def downstream_expansions_set(
     )
 
 
+def ec2_assume_role(
+    role_arn: str,
+    policy: Optional[str] = None,
+    duration_seconds: Optional[int] = None,
+    command_type: Optional[EvgCommandType] = None,
+) -> BuiltInCommand:
+    """
+    Command to call the aws assumeRole API and returns credentials as expansions.
+
+    :param role_arn: string ARN of the role you want to assume.
+    :param policy: string in JSON format that you want to use as an inline session policy.
+    :param duration_seconds: Int in seconds of how long the returned credentials will be valid.
+    :return: ec2.assume_role command.
+    """
+    return BuiltInCommand(
+        command="ec2.assume_role",
+        params={
+            "role_arn": role_arn,
+            "policy": policy,
+            "duration_seconds": duration_seconds,
+        },
+        type=command_type,
+    )
+
+
 def expansions_update(
     updates: Optional[List[KeyValueParam]] = None,
     file: Optional[str] = None,