Create RIPE ATLAS measurement evenly distributed by country

worldwide_measurement.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

'''worldwide_measurement.py: Creates RIPE ATLAS measurement evenly distributed by country.'''


import os
import sys

import bz2
import collections
import datetime
import json
import random
import requests


def load_probes():
    temporary_path = '/tmp/ripe_atlas'
    os.makedirs(temporary_path, exist_ok = True)

    probes_filename = os.path.join(temporary_path, 'probes.json.bz2')
    if not os.path.exists(probes_filename):
        probes_date = (datetime.datetime.utcnow() - datetime.timedelta(days = 1)).date()
        probes_url = 'http://ftp.ripe.net/ripe/atlas/probes/archive/{0:%Y}/{0:%m}/{0:%Y%m%d.json.bz2}'.format(probes_date)
        probes_compressed_data = requests.get(probes_url)
        with open(probes_filename, 'wb') as probes_file:
            probes_file.write(probes_compressed_data.content)

    with bz2.open(probes_filename, mode = 'rt') as probes_data:
        probes_data = json.load(probes_data)

    probes = probes_data['objects']
    return probes


def select_random_probes_by_country(probes):
    probes_by_country = collections.defaultdict(list)
    for probe in probes:
        if probe['status_name'] == 'Connected':
            probes_by_country[probe['country_code']].append(probe)

    selected_probe_ids = set()
    selected_asns = set()
    selected_asns.add(None)
    while len(selected_probe_ids) < 300:
        selected_probe = random.choice(random.choice(list(probes_by_country.items()))[1])

        selected_asn = selected_probe['asn_v4'] or selected_probe['asn_v6']
        if selected_asn in selected_asns:
            continue
        selected_asns.add(selected_asn)

        selected_probe_ids.add(selected_probe['id'])

    return selected_probe_ids


def create_measurement(measurement, key):
    response = requests.post(
        'https://atlas.ripe.net/api/v2/measurements/?key={}'.format(key),
        headers = {
            'Content-Type': 'application/json',
            'Accept': 'application/json'
        },
        data = json.dumps(measurement)
    )
    assert response.status_code in (200, 201), response
    [measurement_id] = json.loads(response.text)['measurements']
    return int(measurement_id)


def main(arguments):
    '''Main function.'''
    if len(arguments) < 2 or '--help' in arguments:
        print('Creates measurement with specified configuration and 300 randomly selected probes.', file = sys.stderr)
        print('Usage: <measurement.json> [key.txt]', file = sys.stderr)
        return os.EX_USAGE

    probes = load_probes()
    selected_probe_ids = select_random_probes_by_country(probes)

    measurement_filename = arguments[1]
    with open(measurement_filename, 'rt') as measurement_file:
        measurement = json.load(measurement_file)
    measurement['probes'].append({
        'value': ','.join(map(str, selected_probe_ids)),
        'type': 'probes',
        'requested': len(selected_probe_ids)
    })

    key_filename = 'key.txt' if len(arguments) < 3 else arguments[2]
    with open(key_filename, 'rt') as key_file:
        key = key_file.read().rstrip()

    measurement_id = create_measurement(measurement, key)
    print('Measurement {} created.'.format(measurement_id))


if __name__ == "__main__":
    sys.exit(main(sys.argv))