From 7137c65353b251a9027ba9e2d8fef259cac46a44 Mon Sep 17 00:00:00 2001
From: Geoff Wilson <geoff@gr-oss.io>
Date: Thu, 7 Nov 2024 11:09:18 -0500
Subject: [PATCH] Add policy cleanup to kv delete

---
 .gitignore                         |  3 +++
 app/lib/clients/vault/key_value.rb |  1 +
 app/lib/clients/vault/policy.rb    | 10 ++++++++++
 test/lib/clients/vault_test.rb     |  9 ++++++---
 4 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 5ba3862..f59fe59 100644
--- a/.gitignore
+++ b/.gitignore
@@ -31,3 +31,6 @@
 
 # Ignore master key for decrypting credentials and more.
 /config/master.key
+
+# Ignore lock files
+/.lock-*
\ No newline at end of file
diff --git a/app/lib/clients/vault/key_value.rb b/app/lib/clients/vault/key_value.rb
index 9f5b59e..a43c033 100644
--- a/app/lib/clients/vault/key_value.rb
+++ b/app/lib/clients/vault/key_value.rb
@@ -17,6 +17,7 @@ def kv_write(identity, path, data)
       def kv_delete(identity, path)
         verify_policy(identity, policy_path(path))
         client.logical.delete("#{kv_mount}/data/#{path}")
+        remove_policy(identity, policy_path(path))
       end
 
       def configure_kv
diff --git a/app/lib/clients/vault/policy.rb b/app/lib/clients/vault/policy.rb
index 1e6eab3..61dd38a 100644
--- a/app/lib/clients/vault/policy.rb
+++ b/app/lib/clients/vault/policy.rb
@@ -28,6 +28,16 @@ def verify_policy(identity, policy_name)
         end
       end
 
+      def remove_policy(identity, policy_name)
+        sub = identity.sub
+        Domain.with_advisory_lock(sub) do
+          policies, metadata = get_entity_data(sub)
+          policies.reject! { |p| p == policy_name }
+          put_entity(sub, policies, metadata)
+        end
+        client.sys.delete_policy(policy_name)
+      end
+
       private
 
       def create_astral_policy
diff --git a/test/lib/clients/vault_test.rb b/test/lib/clients/vault_test.rb
index f34f5eb..0d80b9c 100644
--- a/test/lib/clients/vault_test.rb
+++ b/test/lib/clients/vault_test.rb
@@ -96,7 +96,7 @@ class VaultTest < ActiveSupport::TestCase
 
     # check policy is created
     entity = @client.read_entity(@identity.sub)
-    assert_equal "kv_policy/#{path}", entity.data[:policies][0]
+    assert_includes entity.data[:policies], "kv_policy/#{path}"
 
     # check kv_read denied to other identity
     alt_identity = Identity.new
@@ -111,8 +111,11 @@ class VaultTest < ActiveSupport::TestCase
     # check kv_delete
     del_secret = @client.kv_delete(@identity, path)
     assert del_secret
-    read_secret = @client.kv_read(@identity, path)
-    assert_nil read_secret
+    # check policy is removed
+    entity = @client.read_entity(@identity.sub)
+    assert_not_includes entity.data[:policies], "kv_policy/#{path}"
+    err = assert_raises { @client.kv_read(@identity, path) }
+    assert_kind_of AuthError, err
   end
 
   test "entity_alias methods" do