download_reports.py

#!/usr/bin/env python3

import argparse
import json
import requests
import os


VERSION = '1.1.0'
API_KEY = ''
API_SECRET = ''
ACCOUNT = ''


def get_azure_subscriptions(token, subaccount, tenant_id):
    """
    https://docs.lacework.net/api/v2/docs/#tag/Configs/paths/~1api~1v2~1Configs~1AzureSubscriptions/get
    /api/v2/Configs/AzureSubscriptions?tenantId={tenantId}
    """
    url = f'https://{ACCOUNT}.lacework.net/api/v2/Configs/AzureSubscriptions?tenantId={tenant_id}'
    headers = {'Content-Type':'application/json', 'Authorization':f'Bearer {token}', 'Account-Name': subaccount}
    results = requests.get(url, headers=headers)
    accounts = []
    try:
        subscriptions = results.json().get('data', [])[0].get('subscriptions', [])
    except:
        return []
    for subscription in subscriptions:
        accounts.append({'tenant_id': tenant_id, 'subscription_id': subscription})
    return accounts 


def get_gcp_projects(token, subaccount, org_id):
    """
    https://docs.lacework.net/api/v2/docs/#tag/Configs/paths/~1api~1v2~1Configs~1GcpProjects/get
    /api/v2/Configs/GcpProjects?orgId={orgId}
    """
    url = f'https://{ACCOUNT}.lacework.net/api/v2/Configs/GcpProjects?orgId={org_id}'
    headers = {'Content-Type':'application/json', 'Authorization':f'Bearer {token}', 'Account-Name': subaccount}
    results = requests.get(url, headers=headers)
    accounts = []
    try:
        projects = results.json().get('data', [])[0].get('projects', [])
    except:
        return []
    for project in projects:
        accounts.append({'org_id': org_id, 'project_id': project})
    return accounts 


def get_csp_accounts(token, subaccount):
    """
    https://docs.lacework.net/api/v2/docs/#tag/CloudAccounts/paths/~1api~1v2~1CloudAccounts/get
    /api/v2/CloudAccounts
    """
    url = f'https://{ACCOUNT}.lacework.net/api/v2/CloudAccounts'
    headers = {'Content-Type':'application/json', 'Authorization':f'Bearer {token}', 'Account-Name': subaccount}
    results = requests.get(url, headers=headers)
    accounts = { 'AwsCfg': [], 'GcpCfg': [], 'AzureCfg': [] }
    try:
        csp_accounts = results.json().get('data', [])
    except:
        return accounts
    for csp_account in csp_accounts:
        if csp_account['type'] == 'AwsCfg':
            accounts['AwsCfg'].append(csp_account['data']['awsAccountId'])
        elif csp_account['type'] == 'AzureCfg':
            subscriptions = get_azure_subscriptions(token, subaccount, csp_account['data']['tenantId'])
            accounts['AzureCfg'] = accounts['AzureCfg'] + subscriptions
        elif csp_account['type'] == 'GcpCfg':
            projects = get_gcp_projects(token, subaccount, csp_account['data']['id'])
            accounts['GcpCfg'] = accounts['GcpCfg'] + projects 
    return accounts


def get_access_token():
    url = f'https://{ACCOUNT}.lacework.net/api/v2/access/tokens'
    headers = {'Content-Type': 'application/json', 'X-LW-UAKS': API_SECRET}
    data = {'keyId': API_KEY, 'expiryTime': 36000}
    results = requests.post(url, headers=headers, data=json.dumps(data))
    token = results.json()['token']
    return token


def get_reports(subaccount, token, report, primary_id, secondary_id=None, format='csv'):
    print(f"Getting {report} report for IDs {primary_id}, {secondary_id} in sub-account {subaccount}")
    headers = {'Content-Type':'application/json', 'Authorization':f'Bearer {token}', 'Account-Name': subaccount}
    url = f'https://{ACCOUNT}.lacework.net/api/v2/Reports?primaryQueryId={primary_id.split()[0]}&format={format}&reportType={report}'
    if secondary_id:
        url = f'{url}&secondaryQueryId={secondary_id.split()[0]}'
    results = requests.get(url, headers=headers)
    return results.content


def save_data(filename, data):
    os.makedirs('reports', mode = 0o755, exist_ok = True)
    with open(f'reports/{filename}', 'wb') as file:
        file.write(data)


def parse_args():
    parser = argparse.ArgumentParser(description=f'Pull Lacework compliance reports across multiple sub-accounts.\nVersion: {VERSION}')
    parser.add_argument('-a', '--aws', action='store_true', help='Download AWS reports')
    parser.add_argument('-z', '--azure', action='store_true', help='Download Azure reports')
    parser.add_argument('-g', '--gcp', action='store_true', help='Download GCP reports')
    parser.add_argument('-c', '--download_csv', action='store_true', help='Download CSV reports')
    parser.add_argument('-p', '--download_pdf', action='store_true', help='Download PDF reports')
    return parser.parse_args()


def get_lw_subaccounts(token):
    """
    https://docs.lacework.net/api/v2/docs/#tag/UserProfile
    /api/v2/UserProfile
    """
    url = f'https://{ACCOUNT}.lacework.net/api/v2/UserProfile'
    headers = {'Content-Type':'application/json', 'Authorization':f'Bearer {token}', 'Account-Name': ACCOUNT}
    results = requests.get(url, headers=headers)
    try:
        subaccounts = results.json().get('data', [])
    except:
        return [] 
    return subaccounts


def main():
    args = parse_args()
    if args.download_csv or args.download_pdf:
        token = get_access_token()
        subaccounts = get_lw_subaccounts(token)
        reports = {}
        for subaccount in subaccounts[0]['accounts']:
            accounts = get_csp_accounts(token, subaccount['accountName'])
            if args.aws:
            	for aws_account in accounts['AwsCfg']:
                    if args.download_csv:
                        data = get_reports(subaccount['accountName'], token, 'AWS_CIS_14', aws_account, None, 'csv')
                        save_data(subaccount['accountName'] + f'_{aws_account}.csv', data)
                    if args.download_pdf:
                        data = get_reports(subaccount['accountName'], token, 'AWS_CIS_14', aws_account, None, 'pdf')
                        save_data(subaccount['accountName'] + f'_{aws_account}.pdf', data)
            if args.gcp:
                for gcp_account in accounts['GcpCfg']:
                   if args.download_csv:
                       data = get_reports(subaccount['accountName'], token, 'GCP_CIS13', gcp_account['org_id'], gcp_account['project_id'], 'csv')
                       save_data(gcp_account['org_id'] + '_' + gcp_account['project_id'] + '.csv', data)
                   if args.download_pdf:
                       data = get_reports(subaccount['accountName'], token, 'GCP_CIS13', gcp_account['org_id'], gcp_account['project_id'], 'pdf')
                       save_data(gcp_account['org_id'] + '_' + gcp_account['project_id'] + '.pdf', data)
            if args.azure:
                for azure_account in accounts['AzureCfg']:
                    if args.download_csv:
                        data = get_reports(subaccount['accountName'], token, 'AZURE_CIS_1_5', azure_account['tenant_id'], azure_account['subscription_id'], 'csv')
                        save_data(azure_account['tenant_id'] + '_' + azure_account['subscription_id'] + '.csv', data)
                    if args.download_pdf:
                        data = get_reports(subaccount['accountName'], token, 'AZURE_CIS_1_5', azure_account['tenant_id'], azure_account['subscription_id'], 'pdf')
                        save_data(azure_account['tenant_id'] + '_' + azure_account['subscription_id'] + '.pdf', data)


if __name__ == "__main__":
    main()