update the straightforward common use cases

ansible-collections · Sep 27, 2022 · 840b85f · 840b85f
1 parent 21bf714
commit 840b85f
Show file tree

Hide file tree

Showing 10 changed files with 21 additions and 38 deletions.
diff --git a/plugins/lookup/hashi_vault.py b/plugins/lookup/hashi_vault.py
@@ -271,11 +271,12 @@ def run(self, terms, variables=None, **kwargs):
             self.client = self.helper.get_vault_client(**client_args)
 
             try:
-                self.authenticator.authenticate(self.client)
+                auth = self.authenticator.authenticate(self.client)
             except (NotImplementedError, HashiVaultValueError) as e:
                 raise AnsibleError(e)
 
-            ret.extend(self.get())
+            with auth:
+                ret.extend(self.get())
 
         return ret
 

diff --git a/plugins/lookup/vault_kv1_get.py b/plugins/lookup/vault_kv1_get.py
@@ -195,11 +195,11 @@ def run(self, terms, variables=None, **kwargs):
 
         try:
             self.authenticator.validate()
-            self.authenticator.authenticate(client)
+            auth = self.authenticator.authenticate(client)
         except (NotImplementedError, HashiVaultValueError) as e:
             raise AnsibleError(e)
 
-        try:
+        with auth:
             for term in terms:
                 try:
                     raw = client.secrets.kv.v1.read_secret(path=term, mount_point=engine_mount_point)
@@ -217,7 +217,5 @@ def run(self, terms, variables=None, **kwargs):
                 data = metadata.pop('data')
 
                 ret.append(dict(raw=raw, data=data, secret=data, metadata=metadata))
-        finally:
-            self.authenticator.logout(client)
 
         return ret
diff --git a/plugins/lookup/vault_kv2_get.py b/plugins/lookup/vault_kv2_get.py
@@ -219,11 +219,11 @@ def run(self, terms, variables=None, **kwargs):
 
         try:
             self.authenticator.validate()
-            self.authenticator.authenticate(client)
+            auth = self.authenticator.authenticate(client)
         except (NotImplementedError, HashiVaultValueError) as e:
             raise AnsibleError(e)
 
-        try:
+        with auth:
             for term in terms:
                 try:
                     raw = client.secrets.kv.v2.read_secret_version(path=term, version=version, mount_point=engine_mount_point)
@@ -242,7 +242,5 @@ def run(self, terms, variables=None, **kwargs):
                 secret = data['data']
 
                 ret.append(dict(raw=raw, data=data, secret=secret, metadata=metadata))
-        finally:
-            self.authenticator.logout(client)
 
         return ret
diff --git a/plugins/lookup/vault_read.py b/plugins/lookup/vault_read.py
@@ -119,11 +119,11 @@ def run(self, terms, variables=None, **kwargs):
 
         try:
             self.authenticator.validate()
-            self.authenticator.authenticate(client)
+            auth = self.authenticator.authenticate(client)
         except (NotImplementedError, HashiVaultValueError) as e:
             raise AnsibleError(e)
 
-        try:
+        with auth:
             for term in terms:
                 try:
                     data = client.read(term)
@@ -134,7 +134,5 @@ def run(self, terms, variables=None, **kwargs):
                     raise AnsibleError("The path '%s' doesn't seem to exist." % term)
 
                 ret.append(data)
-        finally:
-            self.authenticator.logout(client)
 
         return ret
diff --git a/plugins/lookup/vault_write.py b/plugins/lookup/vault_write.py
@@ -157,11 +157,11 @@ def run(self, terms, variables=None, **kwargs):
 
         try:
             self.authenticator.validate()
-            self.authenticator.authenticate(client)
+            auth = self.authenticator.authenticate(client)
         except (NotImplementedError, HashiVaultValueError) as e:
             raise_from(AnsibleError(e), e)
 
-        try:
+        with auth:
             for term in terms:
                 try:
                     response = client.write(path=term, wrap_ttl=wrap_ttl, **data)
@@ -188,7 +188,5 @@ def run(self, terms, variables=None, **kwargs):
                     output = response
 
                 ret.append(output)
-        finally:
-            self.authenticator.logout(client)
 
         return ret
diff --git a/plugins/modules/vault_kv1_get.py b/plugins/modules/vault_kv1_get.py
@@ -165,11 +165,11 @@ def run_module():
 
     try:
         module.authenticator.validate()
-        module.authenticator.authenticate(client)
+        auth = module.authenticator.authenticate(client)
     except (NotImplementedError, HashiVaultValueError) as e:
         module.fail_json(msg=to_native(e), exception=traceback.format_exc())
 
-    try:
+    with auth:
         try:
             raw = client.secrets.kv.v1.read_secret(path=path, mount_point=engine_mount_point)
         except hvac.exceptions.Forbidden as e:
@@ -181,8 +181,6 @@ def run_module():
                 msg = "Invalid or missing path ['%s']."
 
             module.fail_json(msg=msg % (path,), exception=traceback.format_exc())
-    finally:
-        module.authenticator.logout(client)
 
     metadata = raw.copy()
     data = metadata.pop('data')

diff --git a/plugins/modules/vault_kv2_get.py b/plugins/modules/vault_kv2_get.py
@@ -181,11 +181,11 @@ def run_module():
 
     try:
         module.authenticator.validate()
-        module.authenticator.authenticate(client)
+        auth = module.authenticator.authenticate(client)
     except (NotImplementedError, HashiVaultValueError) as e:
         module.fail_json(msg=to_native(e), exception=traceback.format_exc())
 
-    try:
+    with auth:
         try:
             raw = client.secrets.kv.v2.read_secret_version(path=path, version=version, mount_point=engine_mount_point)
         except hvac.exceptions.Forbidden as e:
@@ -195,8 +195,6 @@ def run_module():
                 msg="Invalid or missing path ['%s'] with secret version '%s'. Check the path or secret version." % (path, version or 'latest'),
                 exception=traceback.format_exc()
             )
-    finally:
-        module.authenticator.logout(client)
 
     data = raw['data']
     metadata = data['metadata']

diff --git a/plugins/modules/vault_pki_generate_certificate.py b/plugins/modules/vault_pki_generate_certificate.py
@@ -258,11 +258,11 @@ def run_module():
 
     try:
         module.authenticator.validate()
-        module.authenticator.authenticate(client)
+        auth = module.authenticator.authenticate(client)
     except (NotImplementedError, HashiVaultValueError) as e:
         module.fail_json(msg=to_native(e), exception=traceback.format_exc())
 
-    try:
+    with auth:
         try:
             if module.check_mode:
                 data = {}
@@ -273,8 +273,6 @@ def run_module():
                 )
         except hvac.exceptions.VaultError as e:
             module.fail_json(msg=to_native(e), exception=traceback.format_exc())
-    finally:
-        module.authenticator.logout(client)
 
     # generate_certificate is a write operation which always return a new certificate
     module.exit_json(changed=True, data=data)

diff --git a/plugins/modules/vault_read.py b/plugins/modules/vault_read.py
@@ -107,20 +107,18 @@ def run_module():
 
     try:
         module.authenticator.validate()
-        module.authenticator.authenticate(client)
+        auth = module.authenticator.authenticate(client)
     except (NotImplementedError, HashiVaultValueError) as e:
         module.fail_json(msg=to_native(e), exception=traceback.format_exc())
 
-    try:
+    with auth:
         try:
             data = client.read(path)
         except hvac.exceptions.Forbidden as e:
             module.fail_json(msg="Forbidden: Permission Denied to path '%s'." % path, exception=traceback.format_exc())
 
         if data is None:
             module.fail_json(msg="The path '%s' doesn't seem to exist." % path)
-    finally:
-        module.authenticator.logout(client)
 
     module.exit_json(data=data)
 

diff --git a/plugins/modules/vault_write.py b/plugins/modules/vault_write.py
@@ -142,11 +142,11 @@ def run_module():
 
     try:
         module.authenticator.validate()
-        module.authenticator.authenticate(client)
+        auth = module.authenticator.authenticate(client)
     except (NotImplementedError, HashiVaultValueError) as e:
         module.fail_json(msg=to_native(e), exception=traceback.format_exc())
 
-    try:
+    with auth:
         try:
             if module.check_mode:
                 response = {}
@@ -173,8 +173,6 @@ def run_module():
                 output = response.content
         else:
             output = response
-    finally:
-        module.authenticator.logout(client)
 
     module.exit_json(changed=True, data=output)