Download observations from Stealthwatch Cloud and print them as CSV

download_observations.py

#!/usr/bin/env python3
from argparse import ArgumentParser
from csv import DictWriter
from requests import get
from sys import stdout

LIMIT = 1000


def main(tenant, observation_type, user, key, max_count=10000):
    url = 'https://{}.obsrvbl.com/api/v3/observations/{}/'.format(
        tenant, observation_type
    )
    headers = {
        'Authorization': 'ApiKey {}:{}'.format(user, key),
        'Accept': 'application/json'
    }

    writer = None
    offset = 0
    with stdout as f:
        while True:
            params = {'limit': LIMIT, 'offset': offset, 'order_by': 'newest'}
            response = get(url, headers=headers, params=params)
            objects = response.json()['objects']
            offset += len(objects)

            if writer is None:
                if objects:
                    fieldnames = list(objects[0].keys())
                    writer = DictWriter(f, fieldnames)
                    writer.writeheader()

            writer.writerows(objects)

            if (len(objects) < LIMIT) or (offset > max_count):
                break


if __name__ == '__main__':
    parser = ArgumentParser(
        description=(
            'Download observations from Stealthwatch Cloud and print them to '
            'the screen as CSV.'
        )
    )
    parser.add_argument(
        'tenant',
        type=str,
        help='The tenant ID to use (example for example.obsrvbl.com)'
    )
    parser.add_argument(
        'observation_type',
        type=str,
        help=(
            'Observation type (e.g. persistent_external_server_observation_v2)'
        )
    )
    parser.add_argument(
        'user',
        type=str,
        help='User name'
    )
    parser.add_argument(
        'key',
        type=str,
        help='API Key'
    )
    parser.add_argument(
        '--max-count',
        type=int,
        help='Maximum to fetch'
    )
    args = parser.parse_args()
    main(
        args.tenant,
        args.observation_type,
        args.user,
        args.key,
        max_count=args.max_count,
    )