Allow setting Hashicorp Vault token from File (apache#9644)

Author: aneesh-joseph

airflow/providers/hashicorp/_internal_client/vault_client.py

@@ -67,6 +67,9 @@ class _VaultClient(LoggingMixin):  # pylint: disable=too-many-instance-attribute
     :param token: Authentication token to include in requests sent to Vault
         (for ``token`` and ``github`` auth_type).
     :type token: str
+    :param token_path: path to file containing authentication token to include in requests sent to Vault
+        (for ``token`` and ``github`` auth_type).
+    :type token_path: str
     :param username: Username for Authentication (for ``ldap`` and ``userpass`` auth_types).
     :type username: str
     :param password: Password for Authentication (for ``ldap`` and ``userpass`` auth_types).
@@ -110,6 +113,7 @@ def __init__(  # pylint: disable=too-many-arguments
         mount_point: str = "secret",
         kv_engine_version: Optional[int] = None,
         token: Optional[str] = None,
+        token_path: Optional[str] = None,
         username: Optional[str] = None,
         password: Optional[str] = None,
         key_id: Optional[str] = None,
@@ -134,10 +138,10 @@ def __init__(  # pylint: disable=too-many-arguments
         if auth_type not in VALID_AUTH_TYPES:
             raise VaultError(f"The auth_type is not supported: {auth_type}. "
                              f"It should be one of {VALID_AUTH_TYPES}")
-        if auth_type == "token" and not token:
-            raise VaultError("The 'token' authentication type requires 'token'")
-        if auth_type == "github" and not token:
-            raise VaultError("The 'github' authentication type requires 'token'")
+        if auth_type == "token" and not token and not token_path:
+            raise VaultError("The 'token' authentication type requires 'token' or 'token_path'")
+        if auth_type == "github" and not token and not token_path:
+            raise VaultError("The 'github' authentication type requires 'token' or 'token_path'")
         if auth_type == "approle" and not role_id:
             raise VaultError("The 'approle' authentication type requires 'role_id'")
         if auth_type == "kubernetes":
@@ -161,6 +165,7 @@ def __init__(  # pylint: disable=too-many-arguments
         self.auth_type = auth_type
         self.kwargs = kwargs
         self.token = token
+        self.token_path = token_path
         self.auth_mount_point = auth_mount_point
         self.mount_point = mount_point
         self.username = username
@@ -206,7 +211,7 @@ def client(self) -> hvac.Client:
         elif self.auth_type == "radius":
             self._auth_radius(_client)
         elif self.auth_type == "token":
-            _client.token = self.token
+            self._set_token(_client)
         elif self.auth_type == "userpass":
             self._auth_userpass(_client)
         else:
@@ -307,6 +312,13 @@ def _auth_approle(self, _client: hvac.Client) -> None:
         else:
             _client.auth_approle(role_id=self.role_id, secret_id=self.secret_id)
 
+    def _set_token(self, _client: hvac.Client) -> None:
+        if self.token_path:
+            with open(self.token_path) as f:
+                _client.token = f.read()
+        else:
+            _client.token = self.token
+
     def get_secret(self, secret_path: str, secret_version: Optional[int] = None) -> Optional[dict]:
         """
         Get secret value from the KV engine.

airflow/providers/hashicorp/hooks/vault.py

@@ -89,6 +89,9 @@ class VaultHook(BaseHook):
     :param kubernetes_jwt_path: Path for kubernetes jwt token (for ``kubernetes`` auth_type, default:
         ``/var/run/secrets/kubernetes.io/serviceaccount/token``)
     :type kubernetes_jwt_path: str
+    :param token_path: path to file containing authentication token to include in requests sent to Vault
+        (for ``token`` and ``github`` auth_type).
+    :type token_path: str
     :param gcp_key_path: Path to GCP Credential JSON file (for ``gcp`` auth_type)
            Mutually exclusive with gcp_keyfile_dict
     :type gcp_key_path: str
@@ -114,6 +117,7 @@ def __init__(  # pylint: disable=too-many-arguments
         role_id: Optional[str] = None,
         kubernetes_role: Optional[str] = None,
         kubernetes_jwt_path: Optional[str] = None,
+        token_path: Optional[str] = None,
         gcp_key_path: Optional[str] = None,
         gcp_scopes: Optional[str] = None,
         azure_tenant_id: Optional[str] = None,
@@ -178,6 +182,7 @@ def __init__(  # pylint: disable=too-many-arguments
             mount_point=mount_point,
             kv_engine_version=kv_engine_version,
             token=self.connection.password,
+            token_path=token_path,
             username=self.connection.login,
             password=self.connection.password,
             key_id=self.connection.login,

airflow/providers/hashicorp/secrets/vault.py

@@ -69,6 +69,9 @@ class VaultBackend(BaseSecretsBackend, LoggingMixin):
     :param token: Authentication token to include in requests sent to Vault.
         (for ``token`` and ``github`` auth_type)
     :type token: str
+    :param token_path: path to file containing authentication token to include in requests sent to Vault
+        (for ``token`` and ``github`` auth_type).
+    :type token_path: str
     :param username: Username for Authentication (for ``ldap`` and ``userpass`` auth_type).
     :type username: str
     :param password: Password for Authentication (for ``ldap`` and ``userpass`` auth_type).
@@ -114,6 +117,7 @@ def __init__(  # pylint: disable=too-many-arguments
         mount_point: str = 'secret',
         kv_engine_version: int = 2,
         token: Optional[str] = None,
+        token_path: Optional[str] = None,
         username: Optional[str] = None,
         password: Optional[str] = None,
         key_id: Optional[str] = None,
@@ -143,6 +147,7 @@ def __init__(  # pylint: disable=too-many-arguments
             mount_point=mount_point,
             kv_engine_version=kv_engine_version,
             token=token,
+            token_path=token_path,
             username=username,
             password=password,
             key_id=key_id,

tests/providers/hashicorp/_internal_client/test_vault_client.py

@@ -478,6 +478,21 @@ def test_token(self, mock_hvac):
         self.assertEqual(2, vault_client.kv_engine_version)
         self.assertEqual("secret", vault_client.mount_point)
 
+    @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
+    def test_token_path(self, mock_hvac):
+        mock_client = mock.MagicMock()
+        mock_hvac.Client.return_value = mock_client
+        with open('/tmp/test_token.txt', 'w+') as the_file:
+            the_file.write('s.7AU0I51yv1Q1lxOIg1F3ZRAS')
+        vault_client = _VaultClient(auth_type="token",
+                                    token_path="/tmp/test_token.txt", url="http://localhost:8180")
+        client = vault_client.client
+        mock_hvac.Client.assert_called_with(url='http://localhost:8180')
+        client.is_authenticated.assert_called_with()
+        self.assertEqual("s.7AU0I51yv1Q1lxOIg1F3ZRAS", client.token)
+        self.assertEqual(2, vault_client.kv_engine_version)
+        self.assertEqual("secret", vault_client.mount_point)
+
     @mock.patch("airflow.providers.hashicorp._internal_client.vault_client.hvac")
     def test_default_auth_type(self, mock_hvac):
         mock_client = mock.MagicMock()