From 038f5434d8bc8cfc6c4ec19d169017763ae2334d Mon Sep 17 00:00:00 2001
From: Bob Gendler <boberito@mac.com>
Date: Wed, 11 May 2022 11:33:20 -0400
Subject: [PATCH] appends new custom rules instead of replacing old custom

---
 secure-framework-automapping.py | 44 +++++++++++++++++++++++++++------
 1 file changed, 36 insertions(+), 8 deletions(-)

diff --git a/secure-framework-automapping.py b/secure-framework-automapping.py
index 847b7fc..97efd90 100755
--- a/secure-framework-automapping.py
+++ b/secure-framework-automapping.py
@@ -7,6 +7,7 @@
 from collections import OrderedDict
 import subprocess
 import shutil
+import re
 
 parser = argparse.ArgumentParser(description='Run with no option to generate files to /tmp or point to the mscp directory')
 parser.add_argument("-r", "--repo", default="/tmp/", help="Directory mscp is cloned to", type=str)
@@ -106,23 +107,50 @@
     script_path = path[:-6]
     full_path_mapping = os.path.abspath(path_to_framework_mapping)
     subprocess.call(script_path + "/scripts/generate_mapping.py " + full_path_mapping , shell=True)
-    # original = r'original path where the directory is currently stored\directory name'
     ogpath = script_path + "/build/" + framework + "/rules/"
     rules_dir = os.listdir(ogpath)
     for section in rules_dir:
         original = ogpath + "/" + section
         target = script_path + "/custom/rules/" + section
+    
         if os.path.isdir(target):
-            shutil.rmtree(target)
-        fullpath = os.path.abspath(target)
-        shutil.move(original, target)
+            existing_custom_rules = os.listdir(target)
+            section = str(target).split("/")[-1]
+            matched_rules = []
+            for existing_custom_rule in existing_custom_rules:
+                if existing_custom_rule in os.listdir(ogpath + section):
+                    matched_rules.append(ogpath + section + "/" + existing_custom_rule)
+
+            for new_matched_custom in matched_rules:
+            
+                with open(new_matched_custom) as reader:
+                    new_custom_yam = reader.read()
+            
+                regex = r"(?s)(?<=custom\:).*?(?=tags\:)"
+                ref_match = re.search(regex,new_custom_yam)
+                regex = r"(?s)(?<=tags\:).*?($)"
+                tag_match = re.search(regex,new_custom_yam)
+                
+                with open(target + "/" + str(new_matched_custom).split("/")[-1]) as reader:
+                    existing_custom_yam = reader.read()
+                
+                existing_custom_yam = existing_custom_yam.replace("custom:","custom:{}".format(ref_match.group(0).rstrip()))
+                existing_custom_yam = existing_custom_yam.replace("tags:","tags:{}".format(tag_match.group(0).rstrip()))
+                
+                with open(target + "/" + str(new_matched_custom).split("/")[-1], 'w') as rite:
+                    rite.write(existing_custom_yam)  
+    
+            for rule in os.listdir(ogpath + section):
+                if not os.path.exists(target + "/" + rule):
+                    shutil.copyfile(ogpath + section + "/" + rule, target + "/" + rule)
+
+        elif not os.path.isdir(target):
+            # fullpath = os.path.abspath(target)
+            shutil.move(original, target)        
+
         
     custom_baseline_file = script_path + "/build/" + framework + "/baseline/" + framework.lower() + ".yaml"
     custom_baseline_file = custom_baseline_file.replace(" ","\ ").replace("(","\(").replace(")","\)")
     full_path_baseline = os.path.abspath(custom_baseline_file)
     print(script_path + "/scripts/generate_guidance.py -p -x -s " + full_path_baseline)
     subprocess.call(script_path + "/scripts/generate_guidance.py -p -x -s " + full_path_baseline , shell=True)
-
-
-
-