#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
#  =================================================================
#           #     #                 #     #
#           ##    #   ####   #####  ##    #  ######   #####
#           # #   #  #    #  #    # # #   #  #          #
#           #  #  #  #    #  #    # #  #  #  #####      #
#           #   # #  #    #  #####  #   # #  #          #
#           #    ##  #    #  #   #  #    ##  #          #
#           #     #   ####   #    # #     #  ######     #
#
#        ---   The NorNet Testbed for Multi-Homed Systems  ---
#                        https://www.nntb.no
#  =================================================================
#
#  High-Performance Connectivity Tracer (HiPerConTracer)
#  Copyright (C) 2015-2022 by Thomas Dreibholz
#
#  This program 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.
#
#  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/>.
#
#  Contact: dreibh@simula.no

import os
import sys
import io
import re
import datetime
import bz2
import shutil
import socket
import ipaddress
import psycopg2
import configparser

import GeoIP

import urllib3
import json
import time


# ###### Print log message ##################################################
def log(logstring):
   print('\x1b[32m' + datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S') + ': ' + logstring + '\x1b[0m');


# ###### Abort with error ###################################################
def error(logstring):
   sys.stderr.write(datetime.datetime.now().isoformat() + \
                    ' ===== ERROR: ' + logstring + ' =====\n')
   sys.exit(1)



# ###### Main program #######################################################
if len(sys.argv) < 2:
   error('Usage: ' + sys.argv[0] + ' database_configuration')

configFileName = sys.argv[1]
dbServer       = 'localhost'
dbPort         = 5432
dbUser         = 'importer'
dbPassword     = None
dbName         = 'pingtraceroutedb'


# ====== Get parameters =====================================================
parsedConfigFile = configparser.RawConfigParser()
parsedConfigFile.optionxform = str   # Make it case-sensitive!
try:
   parsedConfigFile.readfp(io.StringIO('[root]\n' + open(configFileName, 'r').read()))
except Exception as e:
    error('Unable to read database configuration file' +  sys.argv[1] + ': ' + str(e))
    sys.exit(1)

for parameterName in parsedConfigFile.options('root'):
   parameterValue = parsedConfigFile.get('root', parameterName)
   if parameterName == 'dbserver':
      dbServer = parameterValue
   elif parameterName == 'dbport':
      dbPort = parameterValue
   elif parameterName == 'dbuser':
      dbUser = parameterValue
   elif parameterName == 'dbpassword':
      dbPassword = parameterValue
   elif parameterName == 'database':
      dbName = parameterValue


# ====== Connect to the database ============================================
try:
   dbConnection = psycopg2.connect(host=str(dbServer), port=str(dbPort),
                                   user=str(dbUser), password=str(dbPassword),
                                   dbname=str(dbName))
   dbConnection.autocommit = False
except Exception as e:
    log('Unable to connect to the database!')
    sys.exit(1)

dbCursor = dbConnection.cursor()


# ====== Add all new addresses to AddressInfo table =========================
log('Adding new addresses from Traceroute hops to AddressInfo table ...')
try:
   dbCursor.execute(
      """INSERT INTO AddressInfo ( IP, TimeStamp )
         SELECT DISTINCT t.HopIP, NOW()
         FROM Traceroute t LEFT JOIN AddressInfo a ON t.HopIP = a.IP
         WHERE a.IP IS NULL;"""
   )
   dbConnection.commit()
except Exception as e:
   log('Update failed: ' + str(e))
   sys.exit(1)

log('Adding new addresses from Traceroute sources to AddressInfo table ...')
try:
   dbCursor.execute(
      """INSERT INTO AddressInfo ( IP, TimeStamp )
         SELECT DISTINCT t.FromIP, NOW()
         FROM Traceroute t LEFT JOIN AddressInfo a ON t.FromIP = a.IP
         WHERE a.IP IS NULL;"""
   )
   dbConnection.commit()
except Exception as e:
   log('Update failed: ' + str(e))
   sys.exit(1)

log('Adding new addresses from Ping sources to AddressInfo table ...')
try:
   dbCursor.execute(
      """INSERT INTO AddressInfo ( IP, TimeStamp )
         SELECT DISTINCT p.FromIP, NOW()
         FROM Ping p LEFT JOIN AddressInfo a ON p.FromIP = a.IP
         WHERE a.IP IS NULL;"""
   )
   dbConnection.commit()
except Exception as e:
   log('Update failed: ' + str(e))
   sys.exit(1)

log('Adding new addresses from Ping destinations to AddressInfo table ...')
try:
   dbCursor.execute(
      """INSERT INTO AddressInfo ( IP, TimeStamp )
         SELECT DISTINCT p.ToIP, NOW()
         FROM Ping p LEFT JOIN AddressInfo a ON p.ToIP = a.IP
         WHERE a.IP IS NULL;"""
   )
   dbConnection.commit()
except Exception as e:
   log('Update failed: ' + str(e))
   sys.exit(1)


# ====== GeoIP Location =====================================================
log('Trying to look up geo-location information ...')

def initGeoIP(database, databaseType = GeoIP.GEOIP_STANDARD):
   gi = None
   try:
      gi = GeoIP.open(database, databaseType)
   except Exception as e:
      sys.stdout.write('GeoIP database ' + database + ' is not available: ' + str(e))
      pass
   return gi

geo4 = initGeoIP("/usr/share/GeoIP/GeoLiteCity.dat")
geo6 = initGeoIP("/usr/share/GeoIP/GeoLiteCityv6.dat")
http = urllib3.PoolManager()


inserts = 0
dbCursor.execute(
   """SELECT IP FROM AddressInfo a WHERE a.Country IS NULL OR a.City IS NULL;"""
)
rows = dbCursor.fetchall()
for row in rows:
   updated  = False
   address  = ipaddress.ip_address(row[0])


   # ====== Try MaxMind GeoIP database ======================================
   # print('Trying MaxMindGeoIP for ' + str(address) + ' ...')
   result = None
   try:
      if address.version == 6:
         result = geo6.record_by_addr_v6(str(address))
      else:
         result = geo4.record_by_addr(str(address))
   except:
      continue

   if result != None:
      country      = result['country_name']
      countryCode  = result['country_code']
      region       = result['region_name']
      city         = result['city']
      postalCode   = result['postal_code']
      organisation = None
      asNumber     = None
      latitude     = str(result['latitude'])
      longitude    = str(result['longitude'])
      updated      = True


   # ====== Try IP-API ======================================================
   if ((result == None) or (city == None) or (city == "")):
      # print('Trying IP-API for ' + str(address) + ' ...')
      url = 'http://ip-api.com/json/' + str(address)
      result = None
      try:
         request = http.request('GET', url)
         result = json.loads(request.data.decode('utf8'))

         try:
            if result['status'] == 'success':
               country      = result['country']
               countryCode  = result['countryCode']
               region       = result['regionName']
               city         = result['city']
               postalCode   = result['zip']
               organisation = result['org']
               asNumber     = None
               m = re.match(r'^AS([0-9]+)', result['as'])
               if m != None:
                  asNumber = str(int(m.group(1)))
               latitude     = str(result['lat'])
               longitude    = str(result['lon'])
               updated      = True
         except Exception as e:
            print("Could not get a useful result: " + str(e))
            print(result)

      except Exception as e:
          print("Location could not be determined automatically: " + str(e))

      time.sleep(60/140)   # max. 150/min are allowed! -> 140 should be fine.


   # ====== Update database =================================================
   if updated == True:
      country      = country.replace("'","''")
      sqlStatement = """UPDATE AddressInfo
                        SET TimeStamp=NOW()
                            ,Country='"""     + country     + """'
                            ,CountryCode='""" + countryCode + """'"""
      if region != None:
         region       = region.replace("'","''")
         sqlStatement = sqlStatement + """,Region='"""       + region       + """'"""
      if city != None:
         city         = city.replace("'","''")
         sqlStatement = sqlStatement + """,City='"""         + city         + """'"""
      if postalCode != None:
         sqlStatement = sqlStatement + """,PostalCode='"""   + postalCode   + """'"""
      if organisation != None:
         organisation = organisation.replace("'","''")
         organisation = organisation[:80]
         sqlStatement = sqlStatement + """,Organisation='""" + organisation + """'"""
      if asNumber != None:
         sqlStatement = sqlStatement + """,ASNumber="""      + asNumber
      if latitude != None:
         sqlStatement = sqlStatement + """,Latitude="""      + latitude
      if longitude != None:
         sqlStatement = sqlStatement + """,Longitude="""     + longitude
      sqlStatement = sqlStatement + """ WHERE IP='""" + str(address) + """'"""

      # print(sqlStatement)
      try:
         dbCursor.execute(sqlStatement)
         dbConnection.commit()
         inserts = inserts + 1
      except Exception as e:
         log('Update failed: ' + str(e))
         sys.exit(1)

if inserts > 0:
   sys.stdout.write('Added ' + str(inserts) + ' entries.\n')


# ====== GeoIP AS Numbers ===================================================
log('Trying to look up AS number information ...')

asnum4 = initGeoIP("/usr/share/GeoIP/GeoIPASNum.dat")
asnum6 = initGeoIP("/usr/share/GeoIP/GeoIPASNumv6.dat")

inserts = 0
dbCursor.execute(
   """SELECT IP FROM AddressInfo a WHERE a.ASNumber IS NULL;"""
)
rows = dbCursor.fetchall()
for row in rows:
   address  = ipaddress.ip_address(row[0])
   result = None
   try:
      if address.version == 6:
         result = asnum6.org_by_addr_v6(str(address))
      else:
         result = asnum4.org_by_addr(str(address))
   except:
      continue

   if result != None:
      m = re.match(r'^AS([0-9]+) (.*)$', result)
      if m != None:
         asNumber     = int(m.group(1))
         organisation = m.group(2)

         organisation = organisation[:80]
         sqlStatement = """UPDATE AddressInfo
                           SET TimeStamp=NOW()
                              ,ASNumber=""" + str(asNumber) + """
                              ,Organisation='""" + organisation + """'
                           WHERE IP='""" + str(address) + """'"""
         print(sqlStatement)
         try:
            dbCursor.execute(sqlStatement)
            dbConnection.commit()
            inserts = inserts + 1
         except Exception as e:
            log('Update failed: ' + str(e))
            sys.exit(1)

if inserts > 0:
   sys.stdout.write('Added ' + str(inserts) + ' entries.\n')


# ====== Get FQDN ===========================================================
log('Trying to resolve missing FQDNs ...')
inserts = 0
dbCursor.execute(
   """SELECT IP FROM AddressInfo a WHERE a.FQDN IS NULL;"""
)
rows = dbCursor.fetchall()
for row in rows:
   address  = row[0]
   hostname = None
   try:
      result = socket.gethostbyaddr(address)
      hostname = result[0]
      # sys.stdout.write('DNS lookup for ' + str(address) + ' -> ' + hostname + '\n')
   except:
      continue

   if hostname != None:
      try:
         dbCursor.execute(
            """UPDATE AddressInfo
               SET TimeStamp=NOW(), FQDN='""" + hostname + """'
               WHERE IP='""" + address + """'"""
         )
         dbConnection.commit()
         inserts = inserts + 1
      except Exception as e:
         log('Update failed: ' + str(e))
         sys.exit(1)

if inserts > 0:
   sys.stdout.write('Added ' + str(inserts) + ' entries.\n')


# ====== Vacuum table========================================================
log('Vacuuming table AddressInfo ...')
try:
   old_isolation_level = dbConnection.isolation_level
   dbConnection.set_isolation_level(0)
   dbCursor.execute("""VACUUM FULL AddressInfo""")
   dbConnection.commit()
   dbConnection.set_isolation_level(old_isolation_level)
except Exception as e:
   log('Optimisation failed: ' + str(e))
   sys.exit(1)
