adds working parallelisation

Author: Skybound1

iamspy/cli.py

@@ -1,4 +1,5 @@
 import logging
+import os
 from typing import List, Optional
 import typer
 from iamspy.model import Model
@@ -82,6 +83,7 @@ def who_can(
     strict_conditions: bool = typer.Option(
         False, help="Whether to require conditions to be passed in for any IAM condition checks"
     ),
+    workers: int = typer.Option(os.cpu_count(), "-w", help="Number of parallel worker processes"),
     model: str = typer.Option(DEFAULT_MODEL_FILENAME, "-f"),
 ):
     """
@@ -91,7 +93,7 @@ def who_can(
     if Path(model).is_file():
         m.load_model(model)
 
-    for x in m.who_can(action, resource, conditions, condition_file, strict_conditions):
+    for x in m.who_can(action, resource, conditions, condition_file, strict_conditions, workers):
         print(x)
 
 
@@ -134,6 +136,7 @@ def who_can_batch_resource(
     strict_conditions: bool = typer.Option(
         False, help="Whether to require conditions to be passed in for any IAM condition checks"
     ),
+    workers: int = typer.Option(os.cpu_count(), "-w", help="Number of parallel worker processes"),
     model: str = typer.Option(DEFAULT_MODEL_FILENAME, "-f"),
 ):
     """
@@ -146,7 +149,7 @@ def who_can_batch_resource(
     with open(resources_file) as fs:
         resources = fs.read().strip().split("\n")
 
-    for source, resource in m.who_can_batch_resource(action, resources, conditions, condition_file, strict_conditions):
+    for source, resource in m.who_can_batch_resource(action, resources, conditions, condition_file, strict_conditions, workers):
         print(f"{source} can perform {action} on {resource}")
 
 

iamspy/model.py

@@ -1,6 +1,8 @@
 from typing import List, Optional, Set, Union, Tuple
 import logging
 import json
+import os
+import multiprocessing as mp
 import z3
 import itertools
 from dataclasses import asdict
@@ -14,6 +16,39 @@
 
 logger = logging.getLogger("iamspy.model")
 
+# Module-level worker state and functions — must be at module level to be picklable
+_worker_model = None
+
+
+def _init_worker(model_json: str):
+    global _worker_model
+    _worker_model = Model()
+    _worker_model._model = DataModel(**json.loads(model_json))
+
+
+def _check_source(args):
+    source, action, resource, conditions, condition_file, strict_conditions = args
+    return source if _worker_model.can_i(
+        source=source,
+        action=action,
+        resource=resource,
+        conditions=conditions,
+        condition_file=condition_file,
+        strict_conditions=strict_conditions,
+    ) else None
+
+
+def _check_source_resource(args):
+    source, action, resource, conditions, condition_file, strict_conditions = args
+    return (source, resource) if _worker_model.can_i(
+        source=source,
+        action=action,
+        resource=resource,
+        conditions=conditions,
+        condition_file=condition_file,
+        strict_conditions=strict_conditions,
+    ) else None
+
 
 class Model:
     def __init__(self):
@@ -340,21 +375,30 @@ def who_can(
         conditions: List[str] = [],
         condition_file: Optional[str] = None,
         strict_conditions: bool = False,
+        workers: int = 1,
     ) -> list[str]:
         """
         Used by the CLI to provide the who-can call.
         """
         possible_accounts = self._check_viable_source_accounts(action, resource)
 
+        sources = set()
         for gaad in self._model.gaads.values():
             if gaad.account not in possible_accounts:
                 logger.debug(f"Skipping {gaad.account} GAAD as not a viable source account for {resource}")
                 continue
             logger.debug(f"Checking identities in {gaad.account} GAAD")
-            sources = set()
             for identity in gaad.RoleDetailList + gaad.UserDetailList:
                 sources.add(identity.Arn)
 
+        if workers > 1:
+            model_json = json.dumps(asdict(self._model), default=json_serial)
+            args = [(s, action, resource, conditions, condition_file, strict_conditions) for s in sources]
+            with mp.Pool(workers, initializer=_init_worker, initargs=(model_json,)) as pool:
+                for result in pool.map(_check_source, args):
+                    if result:
+                        yield result
+        else:
             for source in sources:
                 if self.can_i(
                     source=source,
@@ -398,22 +442,31 @@ def who_can_batch_resource(
         conditions: List[str] = [],
         condition_file: Optional[str] = None,
         strict_conditions: bool = False,
+        workers: int = 1,
     ) -> List[Tuple[str, str]]:
         possible_accounts = set()
 
         for resource in resources:
             accounts = self._check_viable_source_accounts(action, resource)
             possible_accounts.update(accounts)
 
+        sources = set()
         for gaad in self._model.gaads.values():
             if gaad.account not in possible_accounts:
                 logger.debug(f"Skipping {gaad.account} GAAD as not a viable source account for {resource}")
                 continue
             logger.debug(f"Checking identities in {gaad.account} GAAD")
-            sources = set()
             for identity in gaad.RoleDetailList + gaad.UserDetailList:
                 sources.add(identity.Arn)
 
+        if workers > 1:
+            model_json = json.dumps(asdict(self._model), default=json_serial)
+            args = [(s, action, r, conditions, condition_file, strict_conditions) for s, r in itertools.product(sources, resources)]
+            with mp.Pool(workers, initializer=_init_worker, initargs=(model_json,)) as pool:
+                for result in pool.map(_check_source_resource, args):
+                    if result:
+                        yield result
+        else:
             for source, resource in itertools.product(sources, resources):
                 if self.can_i(
                     source=source,