#!/usr/bin/env python3
#
# Copyright VyOS maintainers and contributors <maintainers@vyos.io>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later as
# published by the Free Software Foundation.
#
# This program 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 this program.  If not, see <http://www.gnu.org/licenses/>.
#
# Note: script is used both in opmode and by build hook `40-init-geoip-database.chroot`

import argparse
import sys

from vyos.configquery import ConfigTreeQuery
from vyos.geoip import geoip_download_dbip
from vyos.geoip import geoip_download_maxmind
from vyos.geoip import db_initialise
from vyos.geoip import db_is_initialised
from vyos.geoip import db_import_dbip_ranges
from vyos.geoip import db_import_maxmind_ranges
from vyos.geoip import geoip_update
from vyos.utils.process import run

def get_config(conf):
    return (
        conf.get_config_dict(['firewall', 'global-options', 'geoip'], key_mangling=('-', '_'), get_first_key=True,
                                    no_tag_node_value_mangle=True, with_defaults=True),
        conf.get_config_dict(['firewall'], key_mangling=('-', '_'), get_first_key=True,
                                    no_tag_node_value_mangle=True) if conf.exists(['firewall']) else None,
        conf.get_config_dict(['policy'], key_mangling=('-', '_'), get_first_key=True,
                                    no_tag_node_value_mangle=True) if conf.exists(['policy']) else None,
    )

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--init", help="Initialise", action="store_true")
    parser.add_argument('--download', help="Download", action="store_true")
    args = parser.parse_args()

    if args.init:
        db_initialise()
        db_import_dbip_ranges(delete_file=True)
        sys.exit(0)

    conf = ConfigTreeQuery()

    if not conf.exists(['system', 'name-server']):
        print('There are no system name-servers configured')
        sys.exit(1)

    options, firewall, policy = get_config(conf)

    src_addr = options.get('source_address', '')
    vrf = options.get('vrf', None)

    if args.download or not vrf:
        if options['provider'] == 'db-ip':
            print('Downloading latest DB-IP database...')
            if not geoip_download_dbip(source_address=src_addr):
                print('Failed to download, aborting.')
                sys.exit(1)
        elif options['provider'] == 'maxmind':
            account_id = options['maxmind_account_id']
            license_key = options['maxmind_license_key']
            lite = 'maxmind_lite' in options

            print('Downloading latest MaxMind database...')
            if not geoip_download_maxmind(account_id, license_key, lite, source_address=src_addr):
                print('Failed to download, aborting.')
                sys.exit(1)
        if args.download:
            sys.exit(0)
    elif vrf:
        run(['python3', __file__, '--download'], stdout=None, stderr=None, vrf=vrf)

    if not db_is_initialised():
        db_initialise()

    print('Extracting database...')
    if options['provider'] == 'db-ip':
        if not db_import_dbip_ranges(delete_file=True):
            print('Failed to extract, aborting.')
            sys.exit(1)

    elif options['provider'] == 'maxmind':
        if not db_import_maxmind_ranges(delete_file=True):
            print('Failed to extract, aborting.')
            sys.exit(1)

    if not geoip_update(firewall=firewall, policy=policy):
        sys.exit(1)