From 809d7859df99221b0943d81687fc9a3310937a37 Mon Sep 17 00:00:00 2001
From: Jan Graichen <jgraichen@altimos.de>
Date: Tue, 9 Mar 2021 12:13:50 +0100
Subject: [PATCH] Cleanup strategies from merged pillars

If a merge strategy (__: something) had been present in a pillar that
didn't need merging at all, it wasn't removed. This resulted in the `__`
key being present in the final pillar.

This change deeply scrubs merging strategy keys from pillar data. Due to
this scrubbing, which essentially rebuilds all nested dicts and lists on
merge, the previous deepcopy isn't needed anymore.
---
 CHANGELOG.md                    |  1 +
 salt_tower/pillar/tower.py      | 30 +++++++++++++++++++++---------
 test/pillar/test_tower_tower.py | 27 +++++++++++++++++++++++++++
 3 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a67b383..850c73d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 ### Added
 - Add require option to tower.get
 - Experimental: Support custom context in template injected `render` function
+- Deep scrubbing of merge strategies (`__`) from merged pillar data
 
 ## [1.5.2] - 2021-02-10
 ### Fixed
diff --git a/salt_tower/pillar/tower.py b/salt_tower/pillar/tower.py
index 01d96c4..0b9db9e 100644
--- a/salt_tower/pillar/tower.py
+++ b/salt_tower/pillar/tower.py
@@ -262,7 +262,7 @@ def _compile(
         ctx["pillar"] = self
         ctx["tower"] = self
 
-        def render(path, context=None, renderer='text'):
+        def render(path, context=None, renderer="text"):
             if isinstance(context, dict):
                 context = {**ctx, **context}
             else:
@@ -318,15 +318,27 @@ def get_field(self, key, args, kwargs):
 def _merge(tgt, *objects, strategy="merge-last"):
     for obj in objects:
         if isinstance(tgt, dict):
-            tgt = _merge_dict(tgt, copy.deepcopy(obj), strategy)
+            tgt = _merge_dict(tgt, obj, strategy)
         elif isinstance(tgt, list):
-            tgt = _merge_list(tgt, copy.deepcopy(obj), strategy)
+            tgt = _merge_list(tgt, obj, strategy)
         else:
             raise TypeError(f"Cannot merge {type(tgt)}")
 
     return tgt
 
 
+def _merge_clean(obj):
+    if isinstance(obj, dict):
+        return {k: _merge_clean(v) for k, v in obj.items() if k != "__"}
+
+    if isinstance(obj, list):
+        if obj and isinstance(obj[0], dict) and len(obj[0]) == 1 and "__" in obj[0]:
+            obj = obj[1:]
+        return [_merge_clean(v) for v in obj]
+
+    return obj
+
+
 def _merge_dict(tgt, obj, strategy="merge-last"):
     if not isinstance(obj, dict):
         raise TypeError(f"Cannot merge non-dict type, but is {type(obj)}")
@@ -346,7 +358,7 @@ def _merge_dict(tgt, obj, strategy="merge-last"):
             elif key in tgt and isinstance(tgt[key], list) and isinstance(val, list):
                 _merge_list(tgt[key], val, strategy)
             else:
-                tgt[key] = val
+                tgt[key] = _merge_clean(val)
 
     elif strategy == "merge-first":
         for key, val in obj.items():
@@ -355,11 +367,11 @@ def _merge_dict(tgt, obj, strategy="merge-last"):
             elif key in tgt and isinstance(tgt[key], list) and isinstance(val, list):
                 _merge_list(tgt[key], val, strategy)
             elif key not in tgt:
-                tgt[key] = val
+                tgt[key] = _merge_clean(val)
 
     elif strategy == "overwrite":
         tgt.clear()
-        tgt.update(obj)
+        tgt.update(_merge_clean(obj))
 
     else:
         raise ValueError(f"Unknown strategy: {strategy}")
@@ -380,15 +392,15 @@ def _merge_list(tgt, lst, strategy="merge-last"):
                 tgt.remove(val)
 
     elif strategy == "merge-last":
-        tgt.extend(lst)
+        tgt.extend(_merge_clean(lst))
 
     elif strategy == "merge-first":
-        for val in lst:
+        for val in _merge_clean(lst):
             tgt.insert(0, val)
 
     elif strategy == "overwrite":
         del tgt[:]
-        tgt.extend(lst)
+        tgt.extend(_merge_clean(lst))
 
     else:
         raise ValueError(f"Unknown strategy: {strategy}")
diff --git a/test/pillar/test_tower_tower.py b/test/pillar/test_tower_tower.py
index cfea1d3..5c2ccbc 100644
--- a/test/pillar/test_tower_tower.py
+++ b/test/pillar/test_tower_tower.py
@@ -124,6 +124,24 @@ def test_merge_list_strategy_merge_overwrite(tower):
     assert tgt == ['c']
 
 
+def test_merge_list_remove_nested_doubledash(tower):
+    tgt = {'a': 0}
+    mod = {'a': [{'__': 'overwrite'}, 'a', 'b']}
+
+    tower.merge(tgt, mod)
+
+    assert tgt == {'a': ['a', 'b']}
+
+
+def test_merge_list_remove_nested_doubledash_list(tower):
+    tgt = [0]
+    mod = [[{'__': 'overwrite'}, 1]]
+
+    tower.merge(tgt, mod)
+
+    assert tgt == [0, [1]]
+
+
 def test_merge_dict_strategy_remove(tower):
     tgt = {'a': 0, 'b': 1}
     mod = {'__': 'remove', 'a': None}
@@ -160,6 +178,15 @@ def test_merge_dict_strategy_merge_overwrite(tower):
     assert tgt == {'c': 2}
 
 
+def test_merge_dict_remove_nested_doubledash(tower):
+    tgt = {'a': 0}
+    mod = {'a': {'__': 'overwrite', 'key': 1}}
+
+    tower.merge(tgt, mod)
+
+    assert tgt == {'a': {'key': 1}}
+
+
 def test_format(tower):
     tower.update({'app': {'name': 'MyApp'}})