Merge pull request #54 from firstnevyn/Add_kv2_support

Kv2 support with a specified secret key
voxpupuli · Aug 22, 2022 · dfcf1fe · dfcf1fe
2 parents 373729b + fbca058
commit dfcf1fe
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 5 deletions.
diff --git a/lib/puppet/functions/vault_lookup/lookup.rb b/lib/puppet/functions/vault_lookup/lookup.rb
@@ -5,6 +5,7 @@
     optional_param 'Optional[String]', :vault_cert_path_segment
     optional_param 'String', :vault_cert_role
     optional_param 'String', :vault_namespace
+    optional_param 'String', :vault_key
     return_type 'Sensitive'
   end
 
@@ -14,7 +15,8 @@ def lookup(path,
              vault_url = nil,
              vault_cert_path_segment = nil,
              vault_cert_role = nil,
-             vault_namespace = nil)
+             vault_namespace = nil,
+             vault_key = nil)
     if vault_url.nil?
       Puppet.debug 'No Vault address was set on function, defaulting to value from VAULT_ADDR env value'
       vault_url = ENV['VAULT_ADDR']
@@ -44,7 +46,11 @@ def lookup(path,
                                 vault_namespace)
 
     secret_uri = vault_base_uri + "/v1/#{path.delete_prefix('/')}"
-    data = get_secret(client, secret_uri, token, vault_namespace)
+    data = get_secret(client,
+                      secret_uri,
+                      token,
+                      vault_namespace,
+                      vault_key)
     Puppet::Pops::Types::PSensitiveType::Sensitive.new(data)
   end
 
@@ -58,7 +64,7 @@ def auth_login_body(vault_cert_role)
     end
   end
 
-  def get_secret(client, uri, token, namespace)
+  def get_secret(client, uri, token, namespace, key)
     headers = { 'X-Vault-Token' => token, 'X-Vault-Namespace' => namespace }.delete_if { |_key, value| value.nil? }
     secret_response = client.get(uri,
                                  headers: headers,
@@ -68,7 +74,11 @@ def get_secret(client, uri, token, namespace)
       raise Puppet::Error, append_api_errors(message, secret_response)
     end
     begin
-      JSON.parse(secret_response.body)['data']
+      if key.nil?
+        JSON.parse(secret_response.body)['data']
+      else
+        JSON.parse(secret_response.body)['data']['data'][key]
+      end
     rescue StandardError
       raise Puppet::Error, 'Error parsing json secret data from vault response'
     end

diff --git a/spec/functions/lookup_spec.rb b/spec/functions/lookup_spec.rb
@@ -49,7 +49,7 @@
     end
   end
 
-  it 'logs on, requests a secret using a token, and returns the data wrapped in the Sensitive type' do
+  it 'logs on, requests a kv1 secret using a token, and returns the data wrapped in the Sensitive type' do
     vault_server = MockVault.new
     vault_server.mount('/v1/auth/cert/login', AuthSuccess)
     vault_server.mount('/v1/kv/test', SecretLookupSuccess)
@@ -60,6 +60,18 @@
     end
   end
 
+  it 'logs on, requests a kv2 secret using a token, and returns the data wrapped in the Sensitive type' do
+    custom_auth_segment = 'v1/custom/auth/segment'
+    vault_server = MockVault.new
+    vault_server.mount("/#{custom_auth_segment}/login", AuthSuccess)
+    vault_server.mount('/v1/kv/test', SecretLookupSuccessKV2)
+    vault_server.start_vault do |port|
+      result = function.execute('kv/test', "http://127.0.0.1:#{port}", custom_auth_segment, '', '', 'bar')
+      expect(result).to be_a(Puppet::Pops::Types::PSensitiveType::Sensitive)
+      expect(result.unwrap).to eq('baz')
+    end
+  end
+
   it 'is successful when providing a cert_role while authenticating' do
     vault_server = MockVault.new
     vault_server.mount('/v1/auth/cert/login', AuthSuccessWithRole)

diff --git a/spec/mock_vault_helper.rb b/spec/mock_vault_helper.rb
@@ -31,6 +31,29 @@ module PuppetVaultLookupHelpers
     JSON
                         .freeze
 
+  SECRET_SUCCESS_KV2_DATA = <<~JSON
+    {
+      "request_id": "36b0eadc-890b-1da7-b06d-36bda6a9d94d",
+      "lease_id": "",
+      "lease_duration": 0,
+      "renewable": false,
+      "data": {
+        "data": {
+          "bar": "baz"
+        },
+    "metadata": {
+          "created_time": "2022-07-21T12:38:39.082211175Z",
+          "custom_metadata": null,
+          "deletion_time": "",
+          "destroyed": false,
+          "version": 1
+        }
+      },
+      "warnings": null
+    }
+    JSON
+                            .freeze
+
   AUTH_SUCCESS_DATA = <<~JSON
       {
         "request_id": "03d11bd4-b994-c432-150f-5703a75641d1",
@@ -115,6 +138,14 @@ def do_GET(_request, response)
     end
   end
 
+  class SecretLookupSuccessKV2 < WEBrick::HTTPServlet::AbstractServlet
+    def do_GET(_request, response)
+      response.body = SECRET_SUCCESS_KV2_DATA
+      response.status = 200
+      response.content_type = 'application/json'
+    end
+  end
+
   class SecretLookupWarning < WEBrick::HTTPServlet::AbstractServlet
     def do_GET(_request, response)
       response.body = SECRETS_WARNING_DATA