#!/usr/bin/python3
###############################################################################
#                                                                             #
# libloc - A library to determine the location of someone on the Internet     #
#                                                                             #
# Copyright (C) 2020 IPFire Development Team <info@ipfire.org>                #
#                                                                             #
# This library is free software; you can redistribute it and/or               #
# modify it under the terms of the GNU Lesser General Public                  #
# License as published by the Free Software Foundation; either                #
# version 2.1 of the License, or (at your option) any later version.          #
#                                                                             #
# This library 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           #
# Lesser General Public License for more details.                             #
#                                                                             #
###############################################################################

import argparse
import ipaddress
import logging
import math
import re
import sys

# Load our location module
import location
import location.database
import location.importer
from location.i18n import _

# Initialise logging
log = logging.getLogger("location.importer")
log.propagate = 1

INVALID_ADDRESSES = (
	"0.0.0.0",
	"::/0",
	"0::/0",
)

class CLI(object):
	def parse_cli(self):
		parser = argparse.ArgumentParser(
			description=_("Location Importer Command Line Interface"),
		)
		subparsers = parser.add_subparsers()

		# Global configuration flags
		parser.add_argument("--debug", action="store_true",
			help=_("Enable debug output"))

		# version
		parser.add_argument("--version", action="version",
			version="%(prog)s @VERSION@")

		# Database
		parser.add_argument("--database-host", required=True,
			help=_("Database Hostname"), metavar=_("HOST"))
		parser.add_argument("--database-name", required=True,
			help=_("Database Name"), metavar=_("NAME"))
		parser.add_argument("--database-username", required=True,
			help=_("Database Username"), metavar=_("USERNAME"))
		parser.add_argument("--database-password", required=True,
			help=_("Database Password"), metavar=_("PASSWORD"))

		# Update WHOIS
		update_whois = subparsers.add_parser("update-whois", help=_("Update WHOIS Information"))
		update_whois.set_defaults(func=self.handle_update_whois)

		args = parser.parse_args()

		# Enable debug logging
		if args.debug:
			log.setLevel(logging.DEBUG)

		# Print usage if no action was given
		if not "func" in args:
			parser.print_usage()
			sys.exit(2)

		return args

	def run(self):
		# Parse command line arguments
		args = self.parse_cli()

		# Initialise database
		self.db = self._setup_database(args)

		# Call function
		ret = args.func(args)

		# Return with exit code
		if ret:
			sys.exit(ret)

		# Otherwise just exit
		sys.exit(0)

	def _setup_database(self, ns):
		"""
			Initialise the database
		"""
		# Connect to database
		db = location.database.Connection(
			host=ns.database_host, database=ns.database_name,
			user=ns.database_username, password=ns.database_password,
		)

		with db.transaction():
			db.execute("""
				-- autnums
				CREATE TABLE IF NOT EXISTS autnums(number integer, name text);
				CREATE UNIQUE INDEX IF NOT EXISTS autnums_number ON autnums(number);

				-- networks
				CREATE TABLE IF NOT EXISTS networks(network inet, autnum integer, country text);
				CREATE UNIQUE INDEX IF NOT EXISTS networks_network ON networks(network);
			""")

		return db

	def handle_update_whois(self, ns):
		downloader = location.importer.Downloader()

		# Download all sources
		with self.db.transaction():
			# Create some temporary tables to store parsed data
			self.db.execute("""
				CREATE TEMPORARY TABLE _autnums(number integer, organization text)
					ON COMMIT DROP;
				CREATE UNIQUE INDEX _autnums_number ON _autnums(number);

				CREATE TEMPORARY TABLE _organizations(handle text, name text)
					ON COMMIT DROP;
				CREATE UNIQUE INDEX _organizations_handle ON _organizations(handle);
			""")

			for source in location.importer.WHOIS_SOURCES:
				with downloader.request(source, return_blocks=True) as f:
					for block in f:
						self._parse_block(block)

			self.db.execute("""
				INSERT INTO autnums(number, name)
					SELECT _autnums.number, _organizations.name FROM _autnums
						LEFT JOIN _organizations ON _autnums.organization = _organizations.handle
				ON CONFLICT (number) DO UPDATE SET name = excluded.name;
			""")

		# Download all extended sources
		for source in location.importer.EXTENDED_SOURCES:
			with self.db.transaction():
				# Download data
				with downloader.request(source) as f:
					for line in f:
						self._parse_line(line)

	def _parse_block(self, block):
		# Get first line to find out what type of block this is
		line = block[0]

		# aut-num
		if line.startswith("aut-num:"):
			return self._parse_autnum_block(block)

		# organisation
		elif line.startswith("organisation:"):
			return self._parse_org_block(block)

	def _parse_autnum_block(self, block):
		autnum = {}
		for line in block:
			# Split line
			key, val = split_line(line)

			if key == "aut-num":
				m = re.match(r"^(AS|as)(\d+)", val)
				if m:
					autnum["asn"] = m.group(2)

			elif key == "org":
				autnum[key] = val

		# Skip empty objects
		if not autnum:
			return

		# Insert into database
		self.db.execute("INSERT INTO _autnums(number, organization) \
			VALUES(%s, %s) ON CONFLICT (number) DO UPDATE SET \
				organization = excluded.organization",
			autnum.get("asn"), autnum.get("org"),
		)

	def _parse_org_block(self, block):
		org = {}
		for line in block:
			# Split line
			key, val = split_line(line)

			if key in ("organisation", "org-name"):
				org[key] = val

		# Skip empty objects
		if not org:
			return

		self.db.execute("INSERT INTO _organizations(handle, name) \
			VALUES(%s, %s) ON CONFLICT (handle) DO \
			UPDATE SET name = excluded.name",
			org.get("organisation"), org.get("org-name"),
		)

	def _parse_line(self, line):
		# Skip version line
		if line.startswith("2"):
			return

		# Skip comments
		if line.startswith("#"):
			return

		try:
			registry, country_code, type, line = line.split("|", 3)
		except:
			log.warning("Could not parse line: %s" % line)
			return

		# Skip any lines that are for stats only
		if country_code == "*":
			return

		if type in ("ipv6", "ipv4"):
			return self._parse_ip_line(country_code, type, line)

	def _parse_ip_line(self, country, type, line):
		try:
			address, prefix, date, status, organization = line.split("|")
		except ValueError:
			organization = None

			# Try parsing the line without organization
			try:
				address, prefix, date, status = line.split("|")
			except ValueError:
				log.warning("Unhandled line format: %s" % line)
				return

		# Skip anything that isn't properly assigned
		if not status in ("assigned", "allocated"):
			return

		# Cast prefix into an integer
		try:
			prefix = int(prefix)
		except:
			log.warning("Invalid prefix: %s" % prefix)

		# Fix prefix length for IPv4
		if type == "ipv4":
			prefix = 32 - int(math.log(prefix, 2))

		# Try to parse the address
		try:
			network = ipaddress.ip_network("%s/%s" % (address, prefix), strict=False)
		except ValueError:
			log.warning("Invalid IP address: %s" % address)
			return

		self.db.execute("INSERT INTO networks(network, country) \
			VALUES(%s, %s) ON CONFLICT (network) DO \
			UPDATE SET country = excluded.country",
			"%s" % network, country,
		)


def split_line(line):
	key, colon, val = line.partition(":")

	# Strip any excess space
	key = key.strip()
	val = val.strip()

	return key, val

def main():
	# Run the command line interface
	c = CLI()
	c.run()

main()