#! /usr/bin/env python

# This file is part of IVRE.
# Copyright 2011 - 2014 Pierre LALET <pierre.lalet@cea.fr>
#
# IVRE is free software: you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# IVRE is distributed in the hope that it will be useful, but WITHOUT
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
# or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
# License for more details.
#
# You should have received a copy of the GNU General Public License
# along with IVRE. If not, see <http://www.gnu.org/licenses/>.

"""This tool can be used to manage IP addresses related data, such as
AS number and country information.

"""

import ivre.db
import ivre.geoiputils
import ivre.config

if __name__ == '__main__':
    import sys
    reload(sys)
    sys.setdefaultencoding('utf-8')
    import os
    try:
        import argparse
        parser = argparse.ArgumentParser(
            description=__doc__)
        USING_ARGPARSE = True
    except ImportError:
        import optparse
        parser = optparse.OptionParser(
            description=__doc__)
        parser.parse_args_orig = parser.parse_args

        def my_parse_args():
            res = parser.parse_args_orig()
            res[0].ensure_value('ip', res[1])
            return res[0]
        parser.parse_args = my_parse_args
        parser.add_argument = parser.add_option
        USING_ARGPARSE = False
    TORUN = []
    parser.add_argument('--init', '--purgedb', action='store_true',
                        help='Purge or create and initialize the database.')
    parser.add_argument('--ensure-indexes', action='store_true',
                        help='Create missing indexes (will lock the database).')
    parser.add_argument('--download', action='store_true',
                        help='Fetch all data files.')
    parser.add_argument('--country-csv', metavar='FILE',
                        help='Import FILE into countries database.')
    parser.add_argument('--asnum-csv', metavar='FILE',
                        help='Import FILE into AS database.')
    parser.add_argument('--city-csv', metavar='FILE',
                        help='Import FILE into cities database.')
    parser.add_argument('--location-csv', metavar='FILE',
                        help='Import FILE into locations database.')
    parser.add_argument('--import-all', action='store_true',
                        help='Import all files into databases.')
    parser.add_argument('--dont-feed-ipdata-cols', action='store_true',
                        help='Do not feed data to other purpose-specific'
                        ' DB (this only affects the passive DB for now)')
    parser.add_argument('--quiet', "-q", action='store_true',
                        help='Quiet mode.')
    if USING_ARGPARSE:
        parser.add_argument('ip', nargs='*', metavar='IP',
                            help='Display results for specified IP addresses.')
    args = parser.parse_args()
    if args.init:
        if os.isatty(sys.stdin.fileno()):
            sys.stdout.write(
                'This will remove any country/AS information in your '
                'database. Process ? [y/N] ')
            ans = raw_input()
            if ans.lower() == 'y':
                ivre.db.db.data.init()
    if args.ensure_indexes:
        if os.isatty(sys.stdin.fileno()):
            sys.stdout.write(
                'This will lock your database. Process ? [y/N] ')
            ans = raw_input()
            if ans.lower() == 'y':
                ivre.db.db.data.ensure_indexes()
    if args.download:
        ivre.geoiputils.download_all(verbose=not args.quiet)
    if args.city_csv is not None:
        TORUN.append((ivre.db.db.data.feed_geoip_city,
                      [args.city_csv],
                      {} if args.dont_feed_ipdata_cols else
                      {"feedipdata": [ivre.db.db.passive]}))
    if args.country_csv:
        TORUN.append((ivre.db.db.data.feed_geoip_country,
                      [args.country_csv],
                      {} if args.dont_feed_ipdata_cols else
                      {"feedipdata": [ivre.db.db.passive]}))
    if args.asnum_csv:
        TORUN.append((ivre.db.db.data.feed_geoip_asnum,
                      [args.asnum_csv],
                      {} if args.dont_feed_ipdata_cols else
                      {"feedipdata": [ivre.db.db.passive]}))
    if args.location_csv:
        TORUN.append((ivre.db.db.data.feed_city_location,
                      [args.location_csv], {}))
    if args.import_all:
        for function, fname, kwargs in [
                (ivre.db.db.data.feed_geoip_city,
                 'GeoIPCity-Blocks.csv',
                 {} if args.dont_feed_ipdata_cols else
                 {"feedipdata": [ivre.db.db.passive]}),
                (ivre.db.db.data.feed_geoip_country,
                 'GeoIPCountry.csv',
                 {} if args.dont_feed_ipdata_cols else
                 {"feedipdata": [ivre.db.db.passive]}),
                (ivre.db.db.data.feed_geoip_asnum,
                 'GeoIPASNum.csv',
                 {} if args.dont_feed_ipdata_cols else
                 {"feedipdata": [ivre.db.db.passive]}),
                (ivre.db.db.data.feed_city_location,
                 'GeoIPCity-Location.csv', {}),
        ]:
            TORUN.append((function,
                          [os.path.join(ivre.config.GEOIP_PATH,
                                        fname)],
                          kwargs))
    for r in TORUN:
        r[0](*r[1], **r[2])
    for a in args.ip:
        if a.isdigit():
            a = int(a)
        print a
        for i in [ivre.db.db.data.country_byip(a),
                  ivre.db.db.data.as_byip(a),
                  ivre.db.db.data.location_byip(a)]:
            if i:
                for f in i:
                    print '    ', f, i[f]
