fix

local-app/python-tools/cross-organization/org_runner.py

@@ -18,7 +18,7 @@
     def tqdm(iterable, **kwargs): return iterable
 
 # --- VERSIONING ---
-__version__ = "1.6.3"
+__version__ = "1.6.4"
 
 class OrgTaskRunner:
     def __init__(self, args):
@@ -47,7 +47,7 @@ def get_ou_path(self, org_client, entity_id):
         except: return "Unknown", entity_id
 
     def process_account(self, acc, partition, tasks):
-        """Worker thread logic for cross-account execution."""
+        """Worker thread logic with isolated task data storage."""
         thread_session = boto3.Session(profile_name=self.args.profile, region_name=self.args.region)
         sts, org = thread_session.client('sts'), thread_session.client('organizations')
         acc_id, acc_name = acc['Id'], acc['Name']
@@ -57,7 +57,6 @@ def process_account(self, acc, partition, tasks):
         ou_path, ou_id = self.get_ou_path(org, parents[0]['Id']) if parents else ("Orphaned", "N/A")
         ou_path = ou_path if ou_path else "Root"
 
-        # Metadata including Org and OU details
         account_metadata = {
             "org_id": self.org_id,
             "account_id": acc_id, 
@@ -67,7 +66,8 @@ def process_account(self, acc, partition, tasks):
             "ou_id": ou_id
         }
 
-        account_results = {"metadata": account_metadata, "checks": {}}
+        # Isolation: Map each module name to its specific result data to prevent overwriting
+        account_results = {"metadata": account_metadata, "task_data": {}}
 
         try:
             assumed = sts.assume_role(RoleArn=role_arn, RoleSessionName="OrgRunner")
@@ -77,10 +77,10 @@ def process_account(self, acc, partition, tasks):
                 aws_session_token=assumed['Credentials']['SessionToken'],
                 region_name=self.args.region
             )
-            for t_func in tasks:
+            for mod_name, t_func in tasks:
                 res = t_func(m_sess, acc_id, acc_name, self.args.region)
                 account_results["metadata"]["alias"] = res.get("alias", "N/A")
-                account_results["checks"].update(res.get("data", {}))
+                account_results["task_data"][mod_name] = res.get("data", {})
             return account_results, None
         except Exception as e:
             return None, f"FAILED {acc_name}: {str(e)}"
@@ -91,35 +91,31 @@ def run(self):
         org_client = session.client('organizations')
         partition = session.client('sts').get_caller_identity()['Arn'].split(':')[1]
 
-        # Resolve Organization ID
         try:
             self.org_id = org_client.describe_organization()['Organization']['Id']
-        except Exception: pass
+        except: pass
 
-        # Load tasks & build dynamic metadata
         tasks, check_info = [], []
         if self.args.enable_checks:
             sys.path.append(os.getcwd())
             for m in self.args.enable_checks:
                 mod_name = m.replace('.py', '')
                 module = importlib.import_module(mod_name)
-                tasks.append(getattr(module, 'account_task'))
                 v = getattr(module, '__version__', '?.?.?')
+                tasks.append((mod_name, getattr(module, 'account_task')))
                 check_info.append(f"{mod_name} (v{v})")
 
-        # Gather accounts
         all_accounts = [acc for page in org_client.get_paginator('list_accounts').paginate() 
                         for acc in page['Accounts'] if acc['Status'] == 'ACTIVE']
         all_accounts.sort(key=lambda x: x['Name' if self.args.sort == 'name' else 'Id'].lower())
 
-        # HEADER
         print("-" * 100)
         print(f"AWS ORG TASK RUNNER - v{__version__}")
-        print(f"Organization ID: {self.org_id}")
-        print(f"Target Role:     {self.args.role_name}")
-        print(f"Runners:         {self.args.max_workers}")
-        print(f"Enabled Checks:  {', '.join(check_info) if check_info else 'None'}")
-        print(f"Accounts Found:  {len(all_accounts)}")
+        print(f"Org ID:         {self.org_id}")
+        print(f"Target Role:    {self.args.role_name}")
+        print(f"Runners:        {self.args.max_workers}")
+        print(f"Enabled Checks: {', '.join(check_info)}")
+        print(f"Accounts Found: {len(all_accounts)}")
         print("-" * 100)
 
         with ThreadPoolExecutor(max_workers=self.args.max_workers) as executor:
@@ -133,36 +129,27 @@ def run(self):
         if self.args.output:
             ds = datetime.now().strftime("%Y%m%dT%H%M%S")
 
-            # 1. ACCOUNT BASELINE (JSON/CSV)
-            acc_base = f"audit_results.account.{ds}"
-            with open(f"{acc_base}.json", 'w') as f: 
-                json.dump([r['metadata'] for r in self.full_results], f, indent=2)
-            with open(f"{acc_base}.csv", 'w', newline='') as f:
-                w = csv.DictWriter(f, fieldnames=["org_id", "account_id", "account_name", "alias", "ou_path", "ou_id"])
-                w.writeheader()
-                w.writerows([r['metadata'] for r in self.full_results])
-            self.created_files.extend([f"{acc_base}.json", f"{acc_base}.csv"])
-
-            # 2. CHECK SPECIFIC FILES
-            for cn in [m.replace('.py', '') for m in (self.args.enable_checks or [])]:
-                chk_base = f"audit_results.{cn}.{ds}"
-                # Save CSV (Long Format)
+            for mod_name, _ in tasks:
+                chk_base = f"audit_results.{mod_name}.{ds}"
+                # Save CSV specifically for THIS module's data
                 with open(f"{chk_base}.csv", 'w', newline='') as f:
                     w = csv.writer(f)
                     w.writerow(["org_id", "account_id", "account_alias", "region", "field_name", "field_value"])
                     for res in self.full_results:
-                        for reg, fields in res["checks"].items():
+                        mod_data = res["task_data"].get(mod_name, {})
+                        for reg, fields in mod_data.items():
                             for k, v in fields.items():
                                 w.writerow([self.org_id, res["metadata"]["account_id"], res["metadata"]["alias"], reg, k, v])
-                # Save JSON (Nested Format)
+
+                # Save JSON specifically for THIS module's data
                 with open(f"{chk_base}.json", 'w') as f:
                     json.dump([{
                         "org_id": self.org_id,
                         "account_id": r["metadata"]["account_id"], 
                         "alias": r["metadata"]["alias"],
                         "ou_path": r["metadata"]["ou_path"],
                         "ou_id": r["metadata"]["ou_id"],
-                        "data": r["checks"]
+                        "data": r["task_data"].get(mod_name, {})
                     } for r in self.full_results], f, indent=2)
 
                 self.created_files.extend([f"{chk_base}.csv", f"{chk_base}.json"])
@@ -176,7 +163,5 @@ def run(self):
     p.add_argument("--output", nargs='?', const='DEFAULT')
     p.add_argument("--enable-checks", nargs='+')
     p.add_argument("--max-workers", type=int, default=8)
-    p.add_argument("--profile")
-    p.add_argument("--region", default="us-east-1")
-    p.add_argument("--sort", default="name")
+    p.add_argument("--profile"); p.add_argument("--region", default="us-east-1"); p.add_argument("--sort", default="name")
     OrgTaskRunner(p.parse_args()).run()