diff --git a/airflow/providers/hashicorp/_internal_client/vault_client.py b/airflow/providers/hashicorp/_internal_client/vault_client.py index 5d0ef90afca135..ffc338217a3a9a 100644 --- a/airflow/providers/hashicorp/_internal_client/vault_client.py +++ b/airflow/providers/hashicorp/_internal_client/vault_client.py @@ -74,6 +74,7 @@ class _VaultClient(LoggingMixin): :param key_id: Key ID for Authentication (for ``aws_iam`` and ''azure`` auth_type). :param secret_id: Secret ID for Authentication (for ``approle``, ``aws_iam`` and ``azure`` auth_types). :param role_id: Role ID for Authentication (for ``approle``, ``aws_iam`` auth_types). + :param role_arn: AWS arn role (for ``aws_iam`` auth_type) :param kubernetes_role: Role for Authentication (for ``kubernetes`` auth_type). :param kubernetes_jwt_path: Path for kubernetes jwt token (for ``kubernetes`` auth_type, default: ``/var/run/secrets/kubernetes.io/serviceaccount/token``). @@ -103,6 +104,7 @@ def __init__( password: str | None = None, key_id: str | None = None, secret_id: str | None = None, + role_arn: str | None = None, role_id: str | None = None, kubernetes_role: str | None = None, kubernetes_jwt_path: str | None = "/var/run/secrets/kubernetes.io/serviceaccount/token", @@ -161,6 +163,7 @@ def __init__( self.key_id = key_id self.secret_id = secret_id self.role_id = role_id + self.role_arn = role_arn self.kubernetes_role = kubernetes_role self.kubernetes_jwt_path = kubernetes_jwt_path self.gcp_key_path = gcp_key_path @@ -318,15 +321,36 @@ def _auth_azure(self, _client: hvac.Client) -> None: ) def _auth_aws_iam(self, _client: hvac.Client) -> None: - if self.auth_mount_point: - _client.auth.aws.iam_login( - access_key=self.key_id, - secret_key=self.secret_id, - role=self.role_id, - mount_point=self.auth_mount_point, - ) + if self.key_id and self.secret_id: + auth_args = { + "access_key": self.key_id, + "secret_key": self.secret_id, + "role": self.role_id, + } else: - _client.auth.aws.iam_login(access_key=self.key_id, secret_key=self.secret_id, role=self.role_id) + import boto3 + + if self.role_arn: + sts_client = boto3.client("sts") + credentials = sts_client.assume_role(RoleArn=self.role_arn, RoleSessionName="airflow") + auth_args = { + "access_key": credentials["Credentials"]["AccessKeyId"], + "secret_key": credentials["Credentials"]["SecretAccessKey"], + "session_token": credentials["Credentials"]["SessionToken"], + } + else: + session = boto3.Session() + credentials = session.get_credentials() + auth_args = { + "access_key": credentials.access_key, + "secret_key": credentials.secret_key, + "session_token": credentials.token, + } + + if self.auth_mount_point: + auth_args["mount_point"] = self.auth_mount_point + + _client.auth.aws.iam_login(**auth_args) def _auth_approle(self, _client: hvac.Client) -> None: if self.auth_mount_point: diff --git a/airflow/providers/hashicorp/provider.yaml b/airflow/providers/hashicorp/provider.yaml index 18e19e08537b4a..942b9ebde533e7 100644 --- a/airflow/providers/hashicorp/provider.yaml +++ b/airflow/providers/hashicorp/provider.yaml @@ -73,3 +73,9 @@ connection-types: secrets-backends: - airflow.providers.hashicorp.secrets.vault.VaultBackend + +additional-extras: + - name: boto3 + dependencies: + # Require for AWS assume role authentication + - boto3>=1.33.0 diff --git a/airflow/providers/hashicorp/secrets/vault.py b/airflow/providers/hashicorp/secrets/vault.py index 348992a7f69499..b29ae774612af6 100644 --- a/airflow/providers/hashicorp/secrets/vault.py +++ b/airflow/providers/hashicorp/secrets/vault.py @@ -74,6 +74,7 @@ class VaultBackend(BaseSecretsBackend, LoggingMixin): :param key_id: Key ID for Authentication (for ``aws_iam`` and ''azure`` auth_type). :param secret_id: Secret ID for Authentication (for ``approle``, ``aws_iam`` and ``azure`` auth_types). :param role_id: Role ID for Authentication (for ``approle``, ``aws_iam`` auth_types). + :param role_arn: AWS arn role, :param kubernetes_role: Role for Authentication (for ``kubernetes`` auth_type). :param kubernetes_jwt_path: Path for kubernetes jwt token (for ``kubernetes`` auth_type, default: ``/var/run/secrets/kubernetes.io/serviceaccount/token``). @@ -107,6 +108,7 @@ def __init__( key_id: str | None = None, secret_id: str | None = None, role_id: str | None = None, + role_arn: str | None = None, kubernetes_role: str | None = None, kubernetes_jwt_path: str = "/var/run/secrets/kubernetes.io/serviceaccount/token", gcp_key_path: str | None = None, @@ -147,6 +149,7 @@ def __init__( key_id=key_id, secret_id=secret_id, role_id=role_id, + role_arn=role_arn, kubernetes_role=kubernetes_role, kubernetes_jwt_path=kubernetes_jwt_path, gcp_key_path=gcp_key_path, diff --git a/docs/apache-airflow-providers-hashicorp/secrets-backends/hashicorp-vault.rst b/docs/apache-airflow-providers-hashicorp/secrets-backends/hashicorp-vault.rst index f5b8d4e9be47de..3227b0ef58dead 100644 --- a/docs/apache-airflow-providers-hashicorp/secrets-backends/hashicorp-vault.rst +++ b/docs/apache-airflow-providers-hashicorp/secrets-backends/hashicorp-vault.rst @@ -217,6 +217,18 @@ Add "verify": "absolute path to ca-certificate file" backend = airflow.providers.hashicorp.secrets.vault.VaultBackend backend_kwargs = {"connections_path": "airflow-connections", "variables_path": null, "mount_point": "airflow", "url": "http://127.0.0.1:8200", "verify": "/etc/ssl/certs/ca-certificates"} +Vault authentication with AWS Assume Role STS +""""""""""""""""""""""""""""""""""""""""""""" + +Add parameter "role_arn": "The AWS ARN of the role to assume" + +.. code-block:: ini + + [secrets] + backend = airflow.providers.hashicorp.secrets.vault.VaultBackend + backend_kwargs = {"connections_path": "airflow-connections", "variables_path": null, "mount_point": "airflow", "url": "http://127.0.0.1:8200", "auth_type": "aws_iam", "role_arn": "arn:aws:iam::123456789000:role/hashicorp-aws-iam-role"} + + Using multiple mount points """""""""""""""""""""""""""