#!/usr/bin/python

import urllib2
import urllib
import json
import time
import sys
import time
import getopt
import os
import string

# Default values
family = 4
selection = "area"
num_probes = 50
tcp = False
packets = 1

# Parameters
descr =  "DNS Performances of %s"
data = { "definitions": [
    { "query_argument": "fr", "query_class": "IN", "query_type": "SOA",
      "type": "dns", "is_oneoff": True} ],
         "probes": []}
authfile = "%s/.atlas/auth" % os.environ['HOME']

class CredentialsNotFound(Exception):
    pass

class JsonRequest(urllib2.Request):
    def __init__(self, url):
        urllib2.Request.__init__(self, url)
        self.add_header("Content-Type", "application/json")
        self.add_header("Accept", "application/json")
    
def usage(msg=None):
    print >>sys.stderr, "Usage: %s [-4] [-6] [-l N,N,N...] [-n N] [-t] target" % sys.argv[0]
    if msg is not None:
        print >>sys.stderr, msg

try:
    optlist, args = getopt.getopt (sys.argv[1:], "46l:htp:n:",
                               ["list_probes=", "number=", "packets=", "tcp", "help"])
    for option, value in optlist:
        if option == "--help" or option == "-h":
            usage()
            sys.exit(0)
        elif option == "--list_probes" or option == "-l":
            probes = string.split(value, ',')
            selection = "list"
        elif option =="--tcp" or option == "-t":
            tcp = True
        elif option =="--packets" or option == "-p":
            packets = int(value)
        elif option =="--number" or option == "-n":
            packets = int(value)
        elif option == "-4":
            family = "4"
        elif option == "-6":
            family = "6"
        else:
            # Should never occur, it is trapped by getopt
            print >>sys.stderr, "Unknown option %s" % option
            usage()
            sys.exit(1)
except getopt.error, reason:
    usage(reason)
    sys.exit(1)
if len(args) != 1:
    usage()
    sys.exit(1)
target = args[0]
if tcp:
    data["definitions"][0]["use_tcp"] = True
data["definitions"][0]["af"] = family
data["definitions"][0]["target"] = target
data["definitions"][0]["packets"] = packets
data["definitions"][0]["description"] = descr % target
if selection == "area":
    data["probes"].append({ "requested": num_probes, "type": "area" })
elif selection == "list":
    data["probes"].append({ "requested": len(probes), "type": "probes", "value": string.join(probes,',') })
else:
    print >>sys.stderr, "Internal error, invalide selection \"%s\"" % selection
    sys.exit(1)
if not os.path.exists(authfile):
    raise CredentialsNotFound(authfile)
auth = open(authfile)
key = auth.readline()[:-1]
auth.close()
url = "https://atlas.ripe.net/api/v1/measurement/?key=%s" % key
area = "WW"
data["probes"][0]["value"] = area
json_data = json.dumps(data)
try:
    request = JsonRequest(url)
    conn = urllib2.urlopen(request, json_data)
    results = json.load(conn) # TODO: error handling
    print "%s: measurement #%s" % (area, results["measurements"])
    conn.close()
except urllib2.HTTPError as e:
    print >>sys.stderr, ("Fatal error: %s" % e.read())
    raise
conn.close()