1 files changed, 18 insertions, 7 deletions
diff --git a/lib/leap_cli/config/secrets.rb b/lib/leap_cli/config/secrets.rb
index 366ffd3..184d11d 100644
--- a/lib/leap_cli/config/secrets.rb
+++ b/lib/leap_cli/config/secrets.rb
@@ -13,18 +13,29 @@ module LeapCli; module Config
       @discovered_keys = {}
     end
 
-     # we can't use fetch() or get(), since those already have special meanings
-     def retrieve(key, environment=nil)
-       self.fetch(environment||'default', {})[key.to_s]
-     end
+    # we can't use fetch() or get(), since those already have special meanings
+    def retrieve(key, environment)
+      self.fetch(environment, {})[key.to_s]
+    end
+
+    def set(*args, &block)
+      if block_given?
+        set_with_block(*args, &block)
+      else
+        set_without_block(*args)
+      end
+    end
+
+    def set_without_block(key, value, environment)
+      set_with_block(key, environment) {value}
+    end
 
-    def set(key, value, environment=nil)
-      environment ||= 'default'
+    def set_with_block(key, environment, &block)
       key = key.to_s
       @discovered_keys[environment] ||= {}
       @discovered_keys[environment][key] = true
       self[environment] ||= {}
-      self[environment][key] ||= value
+      self[environment][key] ||= yield
     end
 
     #