#!/usr/bin/env python3

"""ssl_collector.py: First version of SSL loader for SSL model"""
__author__      = "Jan Polisensky"

import dns.resolver
import whois
import datetime
from resolvers.Resolver import Resolver

class Dns_whois(Resolver):

    def __init__(self) -> None:
        """
        ! Constructor of the SSL resolver
        """

        super().__init__()
        self.name = "dns_whois"
        self.timeout = 10 # Default timeout
        self.status = None
        self.dns_resolver = dns.resolver.Resolver()
        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 None:
            self.output['error_description'] = message

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

        return self.output

    def __load_whois_data(self, domain_name):
        whois_record = {}
        data = {'success': None, 'whois_data': None}
        
        
        try:
            test_whois = whois.whois("google.com")
        except Exception as e:
            self.log("Cant initialize whois module")
            self.log("Try: pip3 uninstall whois && pip3 install python-whois")
        
        
        try:
            types = ['registrar', 'creation_date', 'expiration_date', 'dnssec', 'emails']
            w = whois.whois(domain_name)
            i = 0
            for type in types:
                try:
                    whois_record[types[i]] = w[types[i]]
                except:
                    whois_record[types[i]] = None

                i=i+1

            data['success'] = True
            data ['whois_data'] = whois_record
            return data

        except Exception as e:
            data['success'] = False
            data['whois_data'] = e
            self.log(self, str(e))
            return data

    def __load_dns_data(self, domain_name, dns_types=['A', 'AAAA', 'CNAME', 'SOA', 'NS', 'MX', 'TXT']):
        dns_records = {}
        errors = []
        data = {'success': None, 'dns_data': None}
        i = 0
        for type in dns_types:
            result = None
            try:
                result = self.dns_resolver.resolve(domain_name, type)
            except Exception as e:
                dns_records[dns_types[i]] = None
                i=i+1
                errors.append(e)
                continue

            if type == 'A':
                self.ip = result[0]
            dns_records[dns_types[i]] = str(result[0])
            i=i+1

        is_void = 0
        for record in dns_records:
            if dns_records[record] is None:
                is_void += 1
        
        if is_void == len(dns_records):
            data['success'] = False
            data['dns_data'] = {"Description": "All dns records are empty", "Record_errors": errors}
            return data

        data['success'] = True
        data['dns_data'] = dns_records
        return data


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

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

        dns_data = self.__load_dns_data(domain_name)

        if dns_data['success'] is False:
            self.log(dns_data['dns_data'])
            return self.__error_message("Error during dns resolve: " + str(dns_data['dns_data']))


        whois_data = self.__load_whois_data(domain_name)

        if whois_data['success'] is False:
            self.log(whois_data['whois_data'])
            return self.__error_message("Error during whois resolve: " + str(whois_data['whois_data']))

    
        result = {'dns_data': dns_data['dns_data'], 'whois_data': whois_data['whois_data']}

        


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


        return self.output
