#!/usr/bin/python
# Usage: iploc.py </path/to/log> </path/to/output>

#import IP2Location
import IP2Location_cached as IP2Location
import atexit
import fileinput
import re
import struct
import sys
import os
import optparse
import socket
import traceback
import logging

from bsddb3 import db


def load_config(fname):
    d = {}
    f = open(fname)
    for l in f:
        l = l.strip()
        if not l or l[0] == "#":
            continue
        key, val = l.split("=", 1)
        if val[0] in ("'", '"'):
            assert val[0] == val[-1]
            val = val[1:-1]
        d[key] = val
    return d

def get_path(fname):
    if fname[0] == "/":
        return fname
    # emulate shell evaluation of "$PWD"
    if fname.startswith("$PWD/"):
        return fname.replace("$PWD/", mypath + "/")
    return fname

mypath = os.path.dirname(sys.argv[0])
if not mypath:
    mypath = "."
config_name = os.path.join(mypath, "config")
config = None

REVERSE_DNS_DB = db.DB()

temp_ident = ""
temp_user = ""


def get_reverse_dns(ip_address):
    # XXX: this works only with IPv4 addresses (because dnshistory works only with IPv4)
    # The keys in the the reverse DNS db are stored as a char encoded string
    # made from the single octet of the IP address.
    key = socket.inet_aton(ip_address)
    value = REVERSE_DNS_DB.get(key)
    if value is None:
        return None
    struct_sz = struct.calcsize("li")
    assert len(value) > struct_sz
    last_date, num_rec = struct.unpack("li", value[:struct_sz])
    value = value[struct_sz:]

    assert num_rec > 0
    for i in xrange(num_rec):
        # struct module doesn't support variable-length fields,
        # so we need to do manual munging
        date = struct.unpack_from("l", value)
        value = value[struct.calcsize("l"):]
        i = value.index("\0")
        fqdn = value[:i]
        value = value[i:]

    # Return last fqdn, which is the latest
    # XXX: really should do time-based lookups for best precision
    return fqdn


def main(file_names):
    global config
    config = load_config(config_name)

    IP2LocObj = IP2Location.IP2Location()
    IP2LocObj.open(get_path(config["LOC_DB_FILE"]))
    IP2LocObj.load_index()

    REVERSE_DNS_DB.open(get_path(config["DNSHISTORY_DB"]), dbtype=db.DB_BTREE)
    atexit.register(REVERSE_DNS_DB.close)

    fi = fileinput.input([file_names])
    for line in fi:
        locations = []
        match = re.search('\d+\.\d+\.\d+\.\d+\s[^ ]+\s[^ ]+', line)
        if not match:
            log.warn("Unmatched %s:%s: '%s':", fi.filename(), fi.filelineno(), line)
            continue
        else:
            try:
                raw_data = match.group()
                data = raw_data.split()
                ip = get_reverse_dns(data[0]) or data[0]
                ident = data[1]
                user = data[2]
                rec = IP2LocObj.get_all(data[0])
                country_short = rec.country_short
                region = rec.region
                city = rec.city
                isp = rec.isp
            except Exception as e:
                log.exception("Exception while processing %s:%s: '%s':", fi.filename(), fi.filelineno(), line)
                continue

        if country_short != '-':
            locations.append(country_short)

        if region != '-':
            locations.append(region)

        if city != '-':
            locations.append(city)

        temp_ident = ' '.join(locations)

        if temp_ident:
            temp_ident = '"{0}"'.format(temp_ident)
        else:
            temp_ident = ident

        if isp != '-':
            temp_user = '"{0}"'.format(isp)
        else:
            temp_user = user

        temp = ip + ' ' + temp_ident + ' ' + temp_user

        sys.stdout.write(line.replace(raw_data, temp))

        del locations


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    log = logging.getLogger(__file__)
    from optparse import OptionParser
    optparser = OptionParser()
    optparser.add_option("--config",
                      help="Use this config file", metavar="FILE")
    options, args = optparser.parse_args()
    if options.config:
        config_name = options.config
    main(args[0])