x

nir0s · Mar 31, 2017 · 9b72e4e · 9b72e4e
1 parent 991fbbe
commit 9b72e4e
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 21 deletions.
diff --git a/README.md b/README.md
@@ -129,6 +129,7 @@ $ ghost get gcp -j
 }
 
 # Modifying an existing key
+# `--add` can be used to add to a key while modify overwrites it.
 $ ghost put gcp token=my_modified_token --modify
 Stashing key...
 

diff --git a/ghost.py b/ghost.py
@@ -214,7 +214,8 @@ def put(self,
             description='',
             encrypt=True,
             lock=False,
-            key_type='secret'):
+            key_type='secret',
+            add=False):
         """Put a key inside the stash
 
         if key exists and modify true: delete and create
@@ -238,6 +239,10 @@ def put(self,
         same goes for the `uid` which will be generated if it didn't
         previously exist.
 
+        `lock` will lock the key to prevent it from being modified or deleted
+
+        `add` allows to add values to an existing key instead of overwriting.
+
         Returns the id of the key in the database
         """
         def assert_key_is_unlocked(existing_key):
@@ -258,11 +263,17 @@ def assert_value_provided_for_new_key(value, existing_key):
         # TODO: This should be refactored. `_handle_existing_key` deletes
         # the key rather implicitly. It shouldn't do that.
         # `existing_key` will be an empty dict if it doesn't exist
-        existing_key = self._handle_existing_key(name, modify)
+        existing_key = self._handle_existing_key(name, modify or add)
         assert_key_is_unlocked(existing_key)
         assert_value_provided_for_new_key(value, existing_key)
 
         if value:
+            # TODO: fix edge case in which encrypt is false and yet we might
+            # try to add to an existing key. encrypt=false is only used when
+            # `load`ing into a new stash, but someone might use it directly
+            # from the API.
+            if add:
+                value = self._update_existing_key(existing_key, value)
             if encrypt:
                 value = self._encrypt(value)
         else:
@@ -290,7 +301,7 @@ def assert_value_provided_for_new_key(value, existing_key):
 
         audit(
             storage=self._storage.db_path,
-            action='MODIFY' if modify else 'PUT',
+            action='MODIFY' if (modify or add) else 'PUT',
             message=json.dumps(dict(
                 key_name=name,
                 value='HIDDEN',
@@ -302,6 +313,29 @@ def assert_value_provided_for_new_key(value, existing_key):
 
         return key_id
 
+    def _update_existing_key(self, existing_key, value):
+        current_value = self._decrypt(existing_key.get('value')).copy()
+        # We update current_value with value to overwrite
+        # existing values if the user provided overriding values
+        current_value.update(value)
+        return current_value
+
+    def _handle_existing_key(self, key_name, modify):
+        existing_key = self._storage.get(key_name) or {}
+        if existing_key and modify:
+            # TODO: Consider replacing this with self.delete(key_name)
+            if not existing_key['lock']:
+                self._storage.delete(key_name)
+        elif existing_key:
+            raise GhostError(
+                'Key `{0}` already exists. Use the modify flag to overwrite'
+                .format(key_name))
+        elif modify:
+            raise GhostError(
+                "Key `{0}` doesn't exist and therefore cannot be modified"
+                .format(key_name))
+        return existing_key
+
     def get(self, key_name, decrypt=True):
         """Return a key with its parameters if it was found.
         """
@@ -522,22 +556,6 @@ def _decrypt(self, hexified_value):
         value = json.loads(jsonified_value)
         return value
 
-    def _handle_existing_key(self, key_name, modify):
-        existing_key = self._storage.get(key_name) or {}
-        if existing_key and modify:
-            # TODO: Consider replacing this with self.delete(key_name)
-            if not existing_key['lock']:
-                self._storage.delete(key_name)
-        elif existing_key:
-            raise GhostError(
-                'Key `{0}` already exists. Use the modify flag to overwrite'
-                .format(key_name))
-        elif modify:
-            raise GhostError(
-                "Key `{0}` doesn't exist and therefore cannot be modified"
-                .format(key_name))
-        return existing_key
-
     def _assert_valid_stash(self):
         if not self._storage.is_initialized:
             raise GhostError(
@@ -1210,6 +1228,10 @@ def init_stash(stash_path, passphrase, passphrase_size, backend):
               '--modify',
               is_flag=True,
               help='Whether to modify an existing key if it exists')
+@click.option('-a',
+              '--add',
+              is_flag=True,
+              help='Whether to add values to an existing key if it exists')
 @click.option('--lock',
               is_flag=True,
               help='Set the key to be locked, preventing its deletion and '
@@ -1227,6 +1249,7 @@ def put_key(key_name,
             description,
             meta,
             modify,
+            add,
             lock,
             key_type,
             stash,
@@ -1250,7 +1273,8 @@ def put_key(key_name,
             metadata=_build_dict_from_key_value(meta),
             description=description,
             lock=lock,
-            key_type=key_type)
+            key_type=key_type,
+            add=add)
         click.echo('Key stashed successfully')
     except GhostError as ex:
         sys.exit(ex)

diff --git a/tests/test_ghost.py b/tests/test_ghost.py
@@ -957,12 +957,31 @@ def test_put_modify_nonexisting_key(self, test_stash):
             test_stash.put('aws', {'key': 'value'}, modify=True)
         assert "therefore cannot be modified" in str(ex.value)
 
+    def test_put_add_nonexisting_key(self, test_stash):
+        with pytest.raises(ghost.GhostError) as ex:
+            test_stash.put('aws', {'key': 'value'}, modify=True)
+        assert "therefore cannot be modified" in str(ex.value)
+
     def test_put_existing_key_no_modify(self, test_stash):
         test_stash.put('aws', {'key': 'value'})
         with pytest.raises(ghost.GhostError) as ex:
             test_stash.put('aws', {'key': 'value'})
         assert "Use the modify flag to overwrite" in str(ex.value)
 
+    def test_put_add_to_existing_key(self, test_stash):
+        test_stash.put('aws', {'key': 'value'})
+        test_stash.put('aws', {'key2': 'value2'}, add=True)
+        key = test_stash.get('aws')
+        assert key['value'] == {'key': 'value', 'key2': 'value2'}
+        assert_in_log('MODIFY')
+
+    def test_put_add_to_existing_key_overwrite_value(self, test_stash):
+        test_stash.put('aws', {'key': 'value', 'key2': 'value2'})
+        test_stash.put('aws', {'key': 'value2'}, add=True)
+        key = test_stash.get('aws')
+        assert key['value'] == {'key': 'value2', 'key2': 'value2'}
+        assert_in_log('MODIFY')
+
     def test_get(self, test_stash):
         def _test_key(key):
             assert isinstance(key, dict)
@@ -1295,7 +1314,7 @@ def test_put_not_initialized(self):
         assert result.exit_code == 1
         assert 'Stash not initialized' in result.output
 
-    def test_put_no_modify(self, test_cli_stash):
+    def test_put_no_modify_or_add(self, test_cli_stash):
         _invoke('put_key aws key=value')
         result = _invoke('put_key aws key=value')
         assert type(result.exception) == SystemExit
@@ -1310,6 +1329,12 @@ def test_modify_locked(self, test_cli_stash):
         assert 'Key `aws` is locked' in result.output
         assert _invoke('get_key aws key').output.strip() == 'value'
 
+    def test_put_add_nonexisting_key(self, test_cli_stash):
+        result = _invoke('put_key aws key=value --add')
+        assert type(result.exception) == SystemExit
+        assert result.exit_code == 1
+        assert "Key `aws` doesn't exist" in result.output
+
     def test_get(self, test_cli_stash):
         _invoke('put_key aws key=value')
         result = _invoke('get_key aws')