#!/usr/bin/env python3

"""Geo.py: Abstract class for geolocation info resolver"""
__author__      = "Jan Polisensky"


from sre_constants import SUCCESS
from resolvers.Resolver import Resolver

import dns.resolver
import requests
import datetime

class Geo(Resolver):
    def __init__(self) -> None:
        """
        ! Constructor of the Geo resolver 
        """
        super().__init__()
        self.name = "geo"
        self.timeout = 10 # Default timeout
        self.status = None
        self.output = {
            "resolver_name" : self.name,
            "success": None,
            "error_description": None,
            "created": None,
            "data": None
        }


    def __error_message(self, message):

        self.output['success'] = False

        if self.output['error_description'] is not None:
            self.output['error_description'] = message

        self.output['created'] = datetime.datetime.now()

        return self.output


    def resolve(self, domain_name, ip_list, timeout=10):

        # array for results
        results = []
        self.timeout = timeout

        if type(domain_name) is not str:
            return self.__error_message("Domain name is not string, got:" + str(type(domain_name)))

        # If the ip_list is empty do manual resolve
        if len(ip_list) < 1:
            resolved_ip = self.__resolve_domain(domain_name)
            if not resolved_ip:
                return self.__error_message("Cant resolve domain name to IP" + str(domain_name))

            ip_list.append(resolved_ip)



        # Geo resolver returns results for each ip adress
        for ip in ip_list:
            ip_result = self.__get_geo_data(ip)
            if not ip_result:
                self.__error_message("Cant get geo data" + str(domain_name))
            results.append(ip_result)


        # Return formated data #
        self.output['success'] = True
        self.output['created'] = datetime.datetime.now()
        self.output['data'] = results


        return self.output


    # Resolver domain to ip adress
    def __resolve_domain(self, domain_name) -> str:
        self.dns_resolver = dns.resolver.Resolver()

        try:
            answer = self.dns_resolver.resolve(domain_name)
            return str(answer[0].to_text())


        except Exception as e:
            self.output['error_description'] = "Error during name resolve:" + str(e)
            return False


    def __get_geo_data(self, ip):

        ### You can try another api key in case of query count limitation #

        #ip_auth_token = "6b3b15bcf578ec"
        ip_auth_token = "f6157341b9e078"
        # ip_auth_token = "7b7427498417ed"

        ####################################################################
        geo_data = {}
        keys = ['country', 'region' ,'city' ,'loc' ,'org']
        url =  "https://ipinfo.io/" + str(ip) + "/?token=" + ip_auth_token

        raw_json = None

        try:
            raw_json = requests.get(url).json()

        except Exception as e:
            self.output['error_description'] = "No data from geo-api, check for API key query count"
            return False

        for i in range(len(keys)):
            try:
                geo_data[keys[i]] = raw_json[keys[i]]
            except:
                geo_data[keys[i]] = None

        # Append IP to identify what IP is the data related to
        geo_data['ip'] = ip

        return geo_data
