async def process_task(task,
add_task) -> AsyncGenerator[Tuple[str, dict], None]:
account_arn = f'arn:aws:iam::{task.account_id}:role/{AUDIT_READER_ROLE}'
account_info = {'account_id': task.account_id}
client_name, method_name = task.method.split('.', 1)
try:
session = _SESSION_CACHE[account_arn] = (
_SESSION_CACHE[account_arn]
if account_arn in _SESSION_CACHE else await aio_sts_assume_role(
src_role_arn=AUDIT_ASSUMER_ARN,
dest_role_arn=account_arn,
dest_external_id=READER_EID,
))
async with session.client(client_name) as client:
if hasattr(client, 'describe_regions'):
response = await client.describe_regions()
region_names = [
region['RegionName'] for region in response['Regions']
]
else:
region_names = API_METHOD_SPECS[task.method].get(
'regions', [None])
for rn in region_names:
async with session.client(client_name, region_name=rn) as client:
async for response in load_task_response(client, task):
if type(response) is DBEntry:
if rn is not None:
response.entity['region'] = rn
yield (task.method, response.entity)
elif type(response) is CollectTask:
add_task(response)
else:
log.info('log response', response)
except ClientError as e:
# record missing auditor role as empty account summary
log.error(e, 'failed processing task')
yield (
task.method,
updated(
account_info,
recorded_at=parse_date(
e.response['ResponseMetadata']['HTTPHeaders']['date']),
),
)
async def fetch(session, url, fetch_over=0) -> dict:
if fetch_over:
await asyncio.sleep(fetch_over * random())
async with session.get(f'https://snowflake.jamfcloud.com/JSSResource{url}',
headers=HEADERS) as response:
txt = await response.text()
date_header = response.headers.get('Date')
if date_header is None:
log.info(f'GET {url} -> status({response.status}) text({txt})')
return {}
result = {'recorded_at': parse_date(date_header)}
try:
return updated(result, json.loads(txt))
except JSONDecodeError:
log.info(f'GET {url} -> status({response.status}) text({txt})')
return result
async def main(table_name):
async with aiohttp.ClientSession() as session:
cids = [
c['id']
for c in (await fetch(session, '/computers')).get('computers', [])
]
log.info(f'loading {len(cids)} computer details')
computers = await asyncio.gather(
*[fetch_computer(session, cid) for cid in cids])
log.info(f'inserting {len(computers)} computers into {table_name}')
rows = [
updated(c.get('computer'),
computer_id=cid,
recorded_at=c.get('recorded_at'))
for cid, c in zip(cids, computers)
]
db.insert(table_name, rows)
return len(rows)
def process_aws_response(task, page):
response_coldict = API_METHOD_SPECS[task.method]['response']
children_list = API_METHOD_SPECS[task.method].get('children', [])
params = API_METHOD_SPECS[task.method].get('params', {})
base_entity = {}
if task.account_id:
base_entity['account_id'] = task.account_id
base_entity.update({v: task.args[k] for k, v in params.items()})
if isinstance(page, (Exception, type(None))):
base_entity['recorded_at'] = datetime.now()
yield DBEntry(base_entity)
return
metadata = page['ResponseMetadata']
base_entity['recorded_at'] = parse_date(metadata['HTTPHeaders']['date'])
if metadata['HTTPStatusCode'] != 200:
yield DBEntry(base_entity)
return
base_entity.update(process_response_items(response_coldict, page))
iterated_entries = list(process_response_lists(response_coldict, page))
for entry in iterated_entries or [base_entity]:
db_entry = DBEntry(updated(base_entity.copy(), entry))
yield db_entry
for child in children_list:
for method in child.get('methods', [child.get('method')]):
request_args = child.get('args', {})
required_args = child.get('required_args', [])
if any(v not in db_entry.entity
for v in request_args.values()):
continue
if not all(db_entry.entity.get(k) for k in required_args):
continue
args = {k: db_entry.entity[v] for k, v in request_args.items()}
yield CollectTask(task.account_id, method, args)
