#! /usr/bin/env python

"""
This script will select the specified number of bridges at random from
the database. If a single value is provided, then we will select the
same number from each distributor, otherwise it will select the
specified number of bridges from the specified distributor.

The resulting file is JSON with the format:

{
    "IP address:Port" {
        "bridge_fingerprint": <fingerprint>,
        "distributor": <distributor>,
        "transport": <pluggable transport method name>,
    }
}

For example,

{
    "123.234.231.213:32581" {
        "bridge_fingerprint": 3f78668071cb40936a9e0fe0023d75359fae8446,
        "distributor": email,
        "transport": flobfs,
    }
}
"""

from __future__ import print_function

import argparse
import datetime
import json
import math
import os
import sqlite3
import sys

TEST_BRIDGES_DEFAULT_INPUT_FILENAME = "bridge_info_map.json"
TEST_BRIDGES_DEFAULT_OUTPUT_FILENAME = "bridge_info_map.json"
TEST_BRIDGES_TOO_OLD = 3600

def loadExistingJsonMap(mapping):
    """Parse the current mapping file, if it exists"""
    mapping = None
    try:
       mapping_fh = open(mapping, 'r')
       mapping_json = json.load(mapping.fh)
    except IOError:
       mapping_json = json.JSONDecoder().decode('{}')
    except TypeError:
       mapping_json = json.JSONDecoder().decode('{}')
    return mapping_json

def writeJsonFile(bridges, outputfn):
    """Write the bridges information to file"""
    try:
        out = open(outputfn, 'a')
        json.dump([bridges], out)
    except IOError, e:
        print("Failed while opening %s: %s", outputfn, e)
        sys.exit(-1)

def getBridges(counts_list, transports, current_mapping):
    """High-level obtainer of bridges from database"""
    bridges_by_dist = dict()
    now = datetime.datetime.now()
    for dist, count in counts_list.items():
        if dist not in bridges_by_dist.keys():
            bridges_by_dist[dist] = list()
        for transport in transports:
            for i in xrange(count):
                bridge = getRandomBridgeIn(dist, transport)
                bridgemap_key = "%s:%d" % (bridge['address'],
                                           bridge['port'])
                last_seen = datetime.datetime.strptime(
                        bridge['last_seen'], "%Y-%m-%d %H:%M")
                # (now - last_seen) returns datetime.timedelta
                while True:
                    if bridgemap_key not in current_mapping.keys():
                        break
                    if (now - last_seen).total_seconds() > TEST_BRIDGES_TOO_OLD:
                        break
                    if bridge not in bridges_by_dist[dist]:
                        break
                    bridge = getRandomBridgeIn(dist, transport)
                if bridge is None:
                    break
                
                bridges_by_dist[dist].append(bridge)
    return bridges_by_dist

def getRandomBridgeIn(dist, transport):
    bridges = ConnectAndGetAllBridgesFromSqlite3(dist, transport)
    if not bridges:
        fake_bridges = {'hex_key': '3f78668071cb40936a9e0fe0023d75359fae8446',
                        'address': '1.2.3.4',
                        'port': 1,
                        'distributor': 'email',
                        'first_seen': '2014-01-01 00:33',
                        'last_seen': '2014-10-26 00:03',
                        }
        return fake_bridges
    bits_needed = math.log(len(bridges), 2)
    index = os.urandom(bits_needed) % len(bridges)
    return bridges[index - 1]

def getBridgesFromSqlite3(dist, transport, cur):
    retBridges = list()
    if not transport:
        cur.execute("SELECT hex_key, address, or_port, distributor, "
                    "first_seen, last_seen FROM Bridges WHERE "
                    "distributor = ?", (distributor, ))
    else:
        # TODO We need to get pluggable transports, too!
        return dict()
    for bridge_info in cur.fetchall():
        bridge = dict(hex_key= bridge_info[0],
                      address=bridge_info[1],
                      port=bridge_info[2],
                      distributor=bridge_info[3],
                      first_seen=bridge_info[4],
                      last_seen=bridge_info[5],
                      transport=transport)
        retBridges.append(bridge)
    return retBridges

def ConnectAndGetAllBridgesFromSqlite3(dist, transport):
    """Establish connection with DB and retrieve bridges"""
    cur = None
    return getBridgesFromSqlite3(dist, transport, cur)

def getDistributorCounts(distributors):
    """Retrieve bridges per distributor from arg"""
    bridges_per_distributors = dict()

    if not distributors:
        return bridges_per_distributors
    for dists in distributors:
        dist = dists.split(":")
        if len(dist) != 2:
            continue
        try:
            count = int(dist[1])
            bridges_per_distributors[dist[0]] = count
        except TypeError:
            print("Fatal: Failed to convert %s to integer", dist[1])
    return bridges_per_distributors

def getOptions():
    """Returns an instance of ArgumentParser with predefined options"""
    description = "Dump to json a set of bridges which will be " \
                  "tested for reachability."
    epilog = "Note: This program is included with BridgeDB, but it " \
             "acts as a helper for retrieving bridges for OONI's " \
             "bridge reachability testing."
    parser = argparse.ArgumentParser(description=description,
                                     epilog=epilog, add_help=True)

    parser.add_argument("--input", "-i",
            default=TEST_BRIDGES_DEFAULT_INPUT_FILENAME)
    parser.add_argument("--output", "-o",
            default=TEST_BRIDGES_DEFAULT_OUTPUT_FILENAME)
    parser.add_argument("--transport", "-t", action="append")
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--count", "-c", action="store_const", const=3)
    group.add_argument("--distributor", "-d", action="append",
            help="Colon separated <distributor>:<number of bridges>")

    return parser

def parseArguments():
    """Returns a Namespace from the parsed args"""
    return getOptions().parse_args()

def main():
    """Kick this thing off"""
    args = parseArguments()
    counts_list = list()
    if args.distributor:
        counts_list = getDistributorCounts(args.distributor)
    else:
        try:
            count = int(args.count)
            counts_list = [count]
        except:
            print("Fatal: Failed to convert %s to integer", args.count)
            sys.exit(-1)

    transports = args.transport
    if not transports:
        transports = [ 'vanilla', ]
    
    current_mapping = loadExistingJsonMap(args.input)
    bridges = getBridges(counts_list, transports, current_mapping)
    writeJsonFile(bridges, args.output)

if __name__ == "__main__":
    main()