Attach scraped data to that from profiles.yaml and output together.

1 files changed, 14 insertions, 14 deletions

diff --git a/scrape_groups_info.py b/scrape_groups_info.py
index 92219db..ac6956b 100644
--- a/scrape_groups_info.py
+++ b/scrape_groups_info.py
@@ -109,28 +109,28 @@ def get_groups_and_techniques(relevant_names):
 
     return groups, all_techniques
 
-def get_group_names(profiles_path):
-    def group_names(inp):
-        return {group["name"] for group in yaml.safe_load(inp)["groups"]}
-
+def get_profiles_data(profiles_path):
     if profiles_path:
         with open(profiles_path) as inp:
-            return group_names(inp)
+            return yaml.safe_load(inp)
 
-    return group_names(sys.stdin)
+    return yaml.safe_load(sys.stdin)
 
 if __name__ == "__main__":
-    group_names = get_group_names(None if len(sys.argv) < 2 else sys.argv[1])
-    groups, techniques = get_groups_and_techniques(group_names)
-    missing_names = group_names.difference(groups)
+    profiles_data = get_profiles_data(None if len(sys.argv) < 2
+                                      else sys.argv[1])
+    group_profiles = dict((g["name"],g) for g in profiles_data["groups"])
+    groups, techniques = get_groups_and_techniques(group_profiles)
+    missing_names = set(group_profiles).difference(groups)
 
     if missing_names:
         print(f"No data found for group(s): {', '.join(sorted(missing_names))}",
               file=sys.stderr)
 
-    out_obj = {
-        "groups": [g.__dict__ for g in groups.values()],
-        "techniques": [t.__dict__ for t in techniques.values()]
-    }
+    for name, group in groups.items():
+        group_profiles[name].update(group.__dict__)
+
+    profiles_data["groups"] = list(group_profiles.values())
+    profiles_data["techniques"] = [t.__dict__ for t in techniques.values()]
 
-    yaml.safe_dump(out_obj, sys.stdout)
+    yaml.safe_dump(profiles_data, sys.stdout)