atlas_probe.py

import argparse

import itertools
import json
import logging
import os
import random
import time

from typing import Dict, List, Optional, Union

from datetime import datetime
from ripe.atlas.cousteau import (
  Dns,
  Ping,
  Traceroute,
  Sslcert,
  AtlasSource,
  AtlasCreateRequest,
  MeasurementTagger,
  ProbeRequest
)

logging.basicConfig()
LOG = logging.getLogger(__name__)
LOG.setLevel(logging.DEBUG)

ATLAS_API_KEY = os.environ['ATLAS_API_KEY']
REQUEST_PER_SOURCE = 5


def atlas_source_for(asn: Optional[str]=None,
                     country_code: Optional[str]=None,
                     n: int=5,
                     af: int=4):
    """
    Create an AtlasSource based on probes that match criteria.

    Request probes, filtering for specified ASN and country code and their
    status (active).
    """
    params: Dict[str, Union[List, str]] = {
        "status": 1,
        "tags": "system-ipv4-works" if af == 4 else "system-ipv6-works",
    }

    if country_code:
        params['country_code'] = country_code
    if asn:
        params['asn'] = asn

    t0 = time.time_ns()

    pr = ProbeRequest(**params)
    # pr has no len, is not subscriptable -> materialize
    probes = list(pr)

    LOG.info("Loaded %d  probes in %.3f", len(probes), (time.time_ns() - t0)/10**9)

    probe_ids = [p['id'] for p in random.sample(probes, min(n, len(probes)))]

    LOG.info('Selected %d connected probes for %s', len(probe_ids), asn)
    if not probe_ids:
        return None

    return AtlasSource(
        type="probes",
        value=",".join(map(str, probe_ids)),
        requested=len(probe_ids),
    )