def lambda_handler(event: Dict[str, Any], _: Any) -> None:
""" entrypoint"""
config = QueryConfig()
logger = Logger()
logger.info(event=QJLogEvents.InitConfig, config=config)
records = event.get("Records", [])
if not records:
raise Exception("No records found")
if len(records) > 1:
raise Exception(
f"More than one record. BatchSize is probably not 1. event: {event}"
)
body = records[0].get("body")
if body is None:
raise Exception(
f"No record body found. BatchSize is probably not 1. event: {event}"
)
body = json.loads(body)
job = schemas.Job(**body)
logger.info(event=QJLogEvents.InitJob, job=job)
logger.info(event=QJLogEvents.RunQueryStart)
query_result = run_query(job=job, config=config)
logger.info(event=QJLogEvents.RunQueryEnd,
num_results=query_result.get_length())
results: List[schemas.Result] = []
if config.account_id_key not in query_result.query_result_set.fields:
raise Exception(
f"Query results must contain field '{config.account_id_key}'")
for q_r in query_result.to_list():
account_id = q_r[config.account_id_key]
result = schemas.Result(
account_id=account_id,
result={
key: val
for key, val in q_r.items() if key != config.account_id_key
},
)
results.append(result)
graph_spec = schemas.ResultSetGraphSpec(
graph_uris_load_times=query_result.graph_uris_load_times)
result_set = schemas.ResultSetCreate(job=job,
graph_spec=graph_spec,
results=results)
api_key = get_api_key(region_name=config.region)
qj_client = QJAPIClient(host=config.api_host,
port=config.api_port,
api_key=api_key)
logger.info(event=QJLogEvents.CreateResultSetStart)
qj_client.create_result_set(result_set=result_set)
logger.info(event=QJLogEvents.CreateResultSetEnd)
def invoke_lambda(
lambda_name: str, lambda_timeout: int, account_scan_lambda_event: AccountScanLambdaEvent
) -> AccountScanResult:
"""Invoke the AccountScan AWS Lambda function
Args:
lambda_name: name of lambda
lambda_timeout: timeout of the lambda. Used to tell the boto3 lambda client to wait
at least this long for a response before timing out.
account_scan_lambda_event: AccountScanLambdaEvent object to serialize to json and send to the lambda
Returns:
AccountScanResult
Raises:
Exception if there was an error invoking the lambda.
"""
logger = Logger()
account_id = account_scan_lambda_event.account_scan_plan.account_id
with logger.bind(lambda_name=lambda_name, lambda_timeout=lambda_timeout, account_id=account_id):
logger.info(event=AWSLogEvents.RunAccountScanLambdaStart)
boto_config = botocore.config.Config(
read_timeout=lambda_timeout + 10, retries={"max_attempts": 0},
)
session = boto3.Session()
lambda_client = session.client("lambda", config=boto_config)
try:
resp = lambda_client.invoke(
FunctionName=lambda_name, Payload=account_scan_lambda_event.json().encode("utf-8")
)
except Exception as invoke_ex:
error = str(invoke_ex)
logger.info(event=AWSLogEvents.RunAccountScanLambdaError, error=error)
raise Exception(
f"Error while invoking {lambda_name} with event {account_scan_lambda_event.json()}: {error}"
) from invoke_ex
payload: bytes = resp["Payload"].read()
if resp.get("FunctionError", None):
function_error = payload.decode()
logger.info(event=AWSLogEvents.RunAccountScanLambdaError, error=function_error)
raise Exception(
f"Function error in {lambda_name} with event {account_scan_lambda_event.json()}: {function_error}"
)
payload_dict = json.loads(payload)
account_scan_result = AccountScanResult(**payload_dict)
logger.info(event=AWSLogEvents.RunAccountScanLambdaEnd)
return account_scan_result