def lambda_handler(event: Dict[str, Any], context: Any) -> None:
"""Entrypoint"""
root = logging.getLogger()
if root.handlers:
for handler in root.handlers:
root.removeHandler(handler)
account_scan_plan_dict = get_required_lambda_event_var(
event, "account_scan_plan")
account_scan_plan = AccountScanPlan.from_dict(account_scan_plan_dict)
scan_id = get_required_lambda_event_var(event, "scan_id")
artifact_path = get_required_lambda_event_var(event, "artifact_path")
max_svc_scan_threads = get_required_lambda_event_var(
event, "max_svc_scan_threads")
preferred_account_scan_regions = get_required_lambda_event_var(
event, "preferred_account_scan_regions")
scan_sub_accounts = get_required_lambda_event_var(event,
"scan_sub_accounts")
artifact_writer = ArtifactWriter.from_artifact_path(
artifact_path=artifact_path, scan_id=scan_id)
account_scanner = AccountScanner(
account_scan_plan=account_scan_plan,
artifact_writer=artifact_writer,
max_svc_scan_threads=max_svc_scan_threads,
preferred_account_scan_regions=preferred_account_scan_regions,
scan_sub_accounts=scan_sub_accounts,
)
scan_results_dict = account_scanner.scan()
scan_results_str = json.dumps(scan_results_dict, default=json_encoder)
json_results = json.loads(scan_results_str)
return json_results
def scan(
self, account_scan_plan: AccountScanPlan
) -> Generator[AccountScanManifest, None, None]:
"""Scan accounts. Return a list of AccountScanManifest objects.
Args:
account_scan_plan: AccountScanPlan defining this scan op
Yields:
AccountScanManifest objects
"""
num_total_accounts = len(account_scan_plan.account_ids)
account_scan_plans = account_scan_plan.to_batches(
max_accounts=self.config.concurrency.max_accounts_per_thread)
num_account_batches = len(account_scan_plans)
num_threads = min(num_account_batches,
self.config.concurrency.max_account_scan_threads)
logger = Logger()
with logger.bind(
num_total_accounts=num_total_accounts,
num_account_batches=num_account_batches,
muxer=self.__class__.__name__,
num_muxer_threads=num_threads,
):
logger.info(event=AWSLogEvents.MuxerStart)
with ThreadPoolExecutor(max_workers=num_threads) as executor:
processed_accounts = 0
futures = []
for sub_account_scan_plan in account_scan_plans:
account_scan_future = self._schedule_account_scan(
executor, sub_account_scan_plan)
futures.append(account_scan_future)
logger.info(
event=AWSLogEvents.MuxerQueueScan,
account_ids=",".join(
sub_account_scan_plan.account_ids),
)
for future in as_completed(futures):
scan_results_dicts = future.result()
for scan_results_dict in scan_results_dicts:
account_id = scan_results_dict["account_id"]
output_artifact = scan_results_dict["output_artifact"]
account_errors = scan_results_dict["errors"]
api_call_stats = scan_results_dict["api_call_stats"]
artifacts = [output_artifact
] if output_artifact else []
account_scan_result = AccountScanManifest(
account_id=account_id,
artifacts=artifacts,
errors=account_errors,
api_call_stats=api_call_stats,
)
yield account_scan_result
processed_accounts += 1
logger.info(event=AWSLogEvents.MuxerStat,
num_scanned=processed_accounts)
logger.info(event=AWSLogEvents.MuxerEnd)
def _schedule_account_scan(self, executor: ThreadPoolExecutor,
account_scan_plan: AccountScanPlan) -> Future:
"""Schedule a local account scan. Note that we serialize the AccountScanPlan
because boto3 sessions are not thread safe.
Args:
executor: ThreadPoolExecutor to submit scan to
account_scan_plan: AccountScanPlans defining this scan
"""
scan_lambda = lambda: local_account_scan(
scan_id=self.scan_id,
account_scan_plan_dict=account_scan_plan.to_dict(),
config=self.config,
)
return executor.submit(scan_lambda)
def _schedule_account_scan(
self, executor: ThreadPoolExecutor, account_scan_plan: AccountScanPlan
) -> Future:
"""Schedule an account scan by calling the AccountScan lambda with
the proper arguments."""
lambda_event = {
"account_scan_plan": account_scan_plan.to_dict(),
"regions": account_scan_plan.regions,
"json_bucket": self.json_bucket,
"key_prefix": self.key_prefix,
"scan_sub_accounts": self.scan_sub_accounts,
}
return executor.submit(
invoke_lambda,
self.account_scan_lambda_name,
self.account_scan_lambda_timeout,
lambda_event,
)
def _schedule_account_scan(
self, executor: ThreadPoolExecutor, account_scan_plan: AccountScanPlan
) -> Future:
"""Schedule an account scan by calling the AccountScan lambda with
the proper arguments."""
lambda_event = {
"account_scan_plan": account_scan_plan.to_dict(),
"scan_id": self.scan_id,
"artifact_path": self.config.artifact_path,
"max_svc_scan_threads": self.config.concurrency.max_svc_scan_threads,
"preferred_account_scan_regions": self.config.scan.preferred_account_scan_regions,
"scan_sub_accounts": self.config.scan.scan_sub_accounts,
}
return executor.submit(
invoke_lambda,
self.account_scan_lambda_name,
self.account_scan_lambda_timeout,
lambda_event,
)
def local_account_scan(
account_scan_plan_dict: Dict[str, Any], scan_sub_accounts: bool, output_dir: Path
) -> Dict[str, Any]:
"""Scan an account.
Args:
account_scan_plan_dict: AccountScanPlan data defining the scan
scan_sub_accounts: if True, scan subaccounts of any org master accounts
output_dir: output artifats to this Path
"""
artifact_writer = FileArtifactWriter(output_dir=output_dir)
account_scan_plan = AccountScanPlan.from_dict(account_scan_plan_dict=account_scan_plan_dict)
account_scanner = AccountScanner(
account_id=account_scan_plan.account_id,
regions=account_scan_plan.regions,
get_session=account_scan_plan.get_session,
artifact_writer=artifact_writer,
scan_sub_accounts=scan_sub_accounts,
max_svc_threads=DEFAULT_MAX_SVC_THREADS,
)
return account_scanner.scan()
def lambda_handler(event, context):
account_scan_plan_dict = get_required_lambda_event_var(
event, "account_scan_plan")
account_scan_plan = AccountScanPlan.from_dict(account_scan_plan_dict)
json_bucket = get_required_lambda_event_var(event, "json_bucket")
key_prefix = get_required_lambda_event_var(event, "key_prefix")
scan_sub_accounts = get_required_lambda_event_var(event,
"scan_sub_accounts")
artifact_writer = S3ArtifactWriter(bucket=json_bucket,
key_prefix=key_prefix)
account_scanner = AccountScanner(
account_id=account_scan_plan.account_id,
regions=account_scan_plan.regions,
get_session=account_scan_plan.get_session,
artifact_writer=artifact_writer,
scan_sub_accounts=scan_sub_accounts,
max_svc_threads=DEFAULT_MAX_SVC_THREADS,
)
scan_results_dict = account_scanner.scan()
scan_results_str = json.dumps(scan_results_dict, default=json_encoder)
json_results = json.loads(scan_results_str)
return json_results
def local_account_scan(
scan_id: str,
account_scan_plan_dict: Dict[str, Any],
config: Config,
) -> List[Dict[str, Any]]:
"""Scan a set of accounts.
Args:
account_scan_plan_dict: AccountScanPlan defining the scan
config: Config object
"""
artifact_writer = ArtifactWriter.from_artifact_path(
artifact_path=config.artifact_path, scan_id=scan_id)
account_scan_plan = AccountScanPlan.from_dict(
account_scan_plan_dict=account_scan_plan_dict)
account_scanner = AccountScanner(
account_scan_plan=account_scan_plan,
artifact_writer=artifact_writer,
max_svc_scan_threads=config.concurrency.max_svc_scan_threads,
preferred_account_scan_regions=config.scan.
preferred_account_scan_regions,
scan_sub_accounts=config.scan.scan_sub_accounts,
)
return account_scanner.scan()
def run_scan(
muxer: AWSScanMuxer,
config: Config,
artifact_writer: ArtifactWriter,
artifact_reader: ArtifactReader,
) -> Tuple[ScanManifest, GraphSet]:
if config.scan.scan_sub_accounts:
account_ids = get_sub_account_ids(config.scan.accounts,
config.access.accessor)
else:
account_ids = config.scan.accounts
account_scan_plan = AccountScanPlan(account_ids=account_ids,
regions=config.scan.regions,
accessor=config.access.accessor)
logger = Logger()
logger.info(event=AWSLogEvents.ScanAWSAccountsStart)
# now combine account_scan_results and org_details to build a ScanManifest
scanned_accounts: List[str] = []
artifacts: List[str] = []
errors: Dict[str, List[str]] = {}
unscanned_accounts: List[str] = []
stats = MultilevelCounter()
graph_set = None
for account_scan_manifest in muxer.scan(
account_scan_plan=account_scan_plan):
account_id = account_scan_manifest.account_id
if account_scan_manifest.artifacts:
for account_scan_artifact in account_scan_manifest.artifacts:
artifacts.append(account_scan_artifact)
artifact_graph_set_dict = artifact_reader.read_json(
account_scan_artifact)
artifact_graph_set = GraphSet.from_dict(
artifact_graph_set_dict)
if graph_set is None:
graph_set = artifact_graph_set
else:
graph_set.merge(artifact_graph_set)
if account_scan_manifest.errors:
errors[account_id] = account_scan_manifest.errors
unscanned_accounts.append(account_id)
else:
scanned_accounts.append(account_id)
else:
unscanned_accounts.append(account_id)
account_stats = MultilevelCounter.from_dict(
account_scan_manifest.api_call_stats)
stats.merge(account_stats)
if graph_set is None:
raise Exception("BUG: No graph_set generated.")
master_artifact_path = artifact_writer.write_json(name="master",
data=graph_set.to_dict())
logger.info(event=AWSLogEvents.ScanAWSAccountsEnd)
start_time = graph_set.start_time
end_time = graph_set.end_time
scan_manifest = ScanManifest(
scanned_accounts=scanned_accounts,
master_artifact=master_artifact_path,
artifacts=artifacts,
errors=errors,
unscanned_accounts=unscanned_accounts,
api_call_stats=stats.to_dict(),
start_time=start_time,
end_time=end_time,
)
artifact_writer.write_json("manifest", data=scan_manifest.to_dict())
return scan_manifest, graph_set