locate-from-db

#! /usr/bin/python3

# usage: locate-from-db output-dir calibration basemap database \
#                       [batch selector...]

import argparse
import collections
import contextlib
import csv
import datetime
import gzip
import multiprocessing
import os
import pickle
import sys
import time
import zlib
from math import inf as Inf

import psycopg2
import psycopg2.extras

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), 'lib')))
import ageo

_time_0 = time.monotonic()

def progress(message, *args):
    global _time_0
    sys.stderr.write(
        ("{}: " + message + "\n").format(
            datetime.timedelta(seconds = time.monotonic() - _time_0),
            *args))

def warning(message, *args):
    sys.stderr.write(
        ("\t*** " + message + "\n").format(*args))

def load_calibration(cfname):
    try:
        with gzip.open(cfname, "rb") as fp:
            return pickle.load(fp)

    except (OSError, zlib.error, pickle.UnpicklingError) as e:
        sys.stderr.write("unable to load calibration: {}: {}\n"
                         .format(cfname, e))
        sys.exit(1)

def get_batch_list(db, selector):
    cur = db.cursor()
    query = ("SELECT b.id, COUNT(*) FROM batches b, measurements m "
             "WHERE b.id = m.batch AND m.rtt > 0")
    if selector:
        query += "AND (" + selector + ") "
    query += "GROUP BY b.id;"
    cur.execute(query)
    batches = []
    for row in cur:
        if row[1] > 0:
            batches.append(row[0])

    progress("{} non-empty batches selected.", len(batches))
    return batches

Position = collections.namedtuple("Position", ("ipv4", "label", "ilabel", "lon", "lat"))
def get_landmark_positions(db, batches):
    cur = db.cursor()
    cur.execute("SELECT DISTINCT h.ipv4, h.label, h.longitude, h.latitude"
                "  FROM hosts h, measurements m"
                " WHERE m.dst = h.ipv4"
                "   AND m.batch = ANY(%s)",
                (batches,))
    rv = {}
    for row in cur:
        try:
            ilabel = int(row[1].partition('-')[2])
        except:
            ilabel = -1
        pos = Position(row[0], row[1], ilabel, row[2], row[3])
        rv[row[0]] = pos

    return rv

def router_for_addr(addr):
    return addr[:addr.rfind('.')] + '.1'

def retrieve_batch(db, batchid):
    cur = db.cursor(cursor_factory=psycopg2.extras.DictCursor)
    cur.execute("""
        SELECT b.id, b.client_lat, b.client_lon, b.client_addr,
               c.label, c.country, c.asn,
               b.proxied, b.proxy_lat, b.proxy_lon, b.proxy_addr,
               p.label, p.country, p.asn,
               b.annot->>'proxy_label' as proxy_label,
               b.annot->>'proxy_provider' as proxy_provider,
               b.annot->>'proxy_alleged_cc2' as proxy_alleged_cc2
          FROM batches b
     LEFT JOIN hosts c ON b.client_addr = c.ipv4
     LEFT JOIN hosts p ON b.proxy_addr  = p.ipv4
         WHERE b.id = %s""", (batchid,))

    # Copy the metadata into a normal dictionary to avoid problems later
    # when it gets stuffed into a Location annotation.
    metadata_raw = cur.fetchone()
    metadata = {}
    metadata.update(metadata_raw)

    # We don't need a fancy cursor for the next steps.
    cur = db.cursor()

    # There's only ever one source for any given batch (and it is
    # always equal to either client_addr or proxy_addr for that
    # batch).  Throw out all measurements that didn't come back with
    # either errno 0 or 111 (that is, success and ECONNREFUSED).  Also
    # throw out RTT zero, which tends to make the calibration choke.
    # And finally throw out the client and proxy address and
    # 127.0.0.1, which will have anomalously short RTTs (since they
    # never hit the network).
    cur.execute("""
        SELECT dst, rtt FROM measurements
         WHERE batch = %s AND rtt > 0
           AND status IN (0, 111)
           AND dst NOT IN ('127.0.0.1', %s, %s)
         """, (batchid, metadata['client_addr'], metadata['proxy_addr']))
    measurements = collections.defaultdict(list)
    for dst, rtt in cur:
        if 0 <= rtt < 5000:
            measurements[dst].append(rtt)
        else:
            warning("out of range: {} -> {}: {}", batchid, dst, rtt)

    if metadata['proxied']:
        # We need to determine and subtract off the travel time _to_
        # the proxy.
        # There are three ways to do this, in decreasing order of
        # accuracy.  Ideally, we have a measurement of the RTT to the
        # proxy's own router.
        router = router_for_addr(metadata['proxy_addr'])
        if router in measurements:
            adjustment = min(measurements[router]) - 5
            metadata['proxy_rtt_estimation_method'] = 'router'
            metadata['proxy_rtt_estimation_addr'] = router
        else:
            # The client itself may also have been a ping destination.
            # We can't look it up by address, but we can look in the
            # hosts table for the location.
            cur.execute("""
                select ipv4 from hosts
                 where abs(latitude - %s) < 0.01
                   and abs(longitude - %s) < 0.01
            """, (metadata['client_lat'], metadata['client_lon']))

            cdest = None
            adjustment = Inf
            for row in cur:
                addr = row[0]
                if addr in measurements:
                    cadj = min(measurements[addr])
                    if cadj < adjustment:
                        cdest = addr
                        adjustment = cadj

            if cdest is not None:
                # If we have a ping destination that is colocated with
                # the client, then we can estimate the RTT to the proxy
                # as half of the RTT through the proxy and back to the
                # client's location, minus a small fudge factor.
                adjustment = adjustment/2 - 5
                metadata['proxy_rtt_estimation_method'] = 'there_and_back'
                metadata['proxy_rtt_estimation_addr'] = cdest

        # In no case allow the adjustment to be greater than the
        # smallest available measurement minus five milliseconds, nor
        # allow it to be negative.
        cdest = None
        clamp = Inf
        for addr, meas in measurements.items():
            mmeas = min(meas)
            if mmeas < clamp:
                clamp = mmeas
                cdest = addr
        if clamp - 5 < adjustment:
            metadata['proxy_rtt_estimation_clamp'] = clamp
            metadata['proxy_rtt_estimation_clamp_addr'] = cdest
            if 'proxy_rtt_estimation_method' in metadata:
                metadata['proxy_rtt_estimation_method'] += '_clamped'
                metadata['proxy_rtt_estimation_unclamped'] = adjustment
            else:
                metadata['proxy_rtt_estimation_method'] = 'clamp'

            adjustment = clamp - 5
        adjustment = max(adjustment, 0)
        metadata['estimated_proxy_rtt'] = adjustment

        # This loop mutates measurements, so pull the keys up front.
        for addr in list(measurements.keys()):
            measurements[addr] = sorted(
                max(0.1, m - adjustment)
                for m in measurements[addr]
            )

        progress("{}: adjust by {} (method {})",
                 metadata['id'], adjustment,
                 metadata['proxy_rtt_estimation_method'])

    # convert to a normal dict for returning
    return metadata, { k:v for k,v in measurements.items() }

# these are filled in in main()
positions    = None
basemap      = None
calibrations = None

def process_batch(args):
    global positions, basemap
    odir, mode, metadata, measurements = args
    tag, cals, ranging, use_all = mode
    bnd = basemap.bounds
    obsv = []
    for landmark, rtts in measurements.items():
        if landmark not in positions:
            continue
        lpos = positions[landmark]
        if use_all:
            calibration = cals[0]
        elif landmark in cals:
            calibration = cals[landmark]
        elif lpos.label in cals:
            calibration = cals[lpos.label]
        elif lpos.ilabel in cals:
            calibration = cals[lpos.ilabel]
        else:
            continue

        obs = ageo.Observation(
            basemap=basemap,
            ref_lat=lpos.lat,
            ref_lon=lpos.lon,
            range_fn=ranging,
            calibration=calibration,
            rtts=rtts)
        obsv.append(obs)
        bnd = bnd.intersection(obs.bounds)
        if bnd.is_empty:
            return tag + " (empty intersection region)", str(metadata['id'])

    if not obsv:
        return tag + " (no observations)", str(metadata['id'])

    loc = obsv[0]
    for obs in obsv[1:]:
        loc = loc.intersection(obs, bnd)

    #loc = loc.intersection(basemap, bnd)
    loc.annotations.update(metadata)
    loc.save(os.path.join(odir, tag + "-" + str(metadata['id']) + ".h5"))
    return tag, metadata['id']

def marshal_batches(args, db, batches, modes):
    for batchid in batches:
        metadata, measurements = retrieve_batch(db, batchid)
        for mode in modes:
            yield (args.output_dir, mode, metadata, measurements)

def inner_main(args, pool, db, batches):
    global calibrations

    # FIXME: duplicates code from 'calibrate'
    cal_cbg, cal_oct, cal_spo = calibrations
    minmax = ageo.ranging.MinMax
    gaussn = ageo.ranging.Gaussian
    crunch_modes = [
        ("cbg-m-1", cal_cbg, minmax, False),
#       ("cbg-m-a", cal_cbg, minmax, True),
        ("oct-m-1", cal_oct, minmax, False),
#       ("oct-m-a", cal_oct, minmax, True),
#       ("spo-m-1", cal_spo, minmax, False),
        ("spo-m-a", cal_spo, minmax, True),
#       ("spo-g-1", cal_spo, gaussn, False),
        ("spo-g-a", cal_spo, gaussn, True)
    ]

    for tag, id in pool.imap_unordered(
            process_batch,
            marshal_batches(args, db, batches, crunch_modes)):
        progress("{}: {}", id, tag)

def main():
    global positions, calibrations, basemap

    ap = argparse.ArgumentParser()
    ap.add_argument("output_dir")
    ap.add_argument("calibration")
    ap.add_argument("basemap")
    ap.add_argument("database")
    ap.add_argument("batch_selector", nargs=argparse.REMAINDER)
    args = ap.parse_args()

    args.batch_selector = " ".join(args.batch_selector)

    # The worker pool must be created before the database connection,
    # so that the database handle is not duplicated into the worker
    # processes.  However, we also need the database connection before
    # the worker processes exist, to load up 'positions', which is
    # propagated into the worker processes by fork().  So we have to
    # drop the database connection and pick it back up again.
    progress("preparing...")
    os.makedirs(args.output_dir, exist_ok=True)
    calibrations = load_calibration(args.calibration)
    basemap = ageo.Map(args.basemap)
    with contextlib.closing(psycopg2.connect(dbname=args.database)) as db:
        batches = get_batch_list(db, args.batch_selector)
        positions = get_landmark_positions(db, batches)

    with multiprocessing.Pool() as pool:
        with contextlib.closing(psycopg2.connect(dbname=args.database)) as db:
            inner_main(args, pool, db, batches)

main()