#!/usr/bin/env python3

"""DataResolver.py: Daemon for the resolution of external domain-related data"""
__author__      = "Radek Hranicky"

import logging
import json
import time
import concurrent.futures
import copy

from datetime import datetime
from time import sleep
from sqlalchemy import exc
from sqlalchemy.sql.expression import true, false

from database import db
from database.models import OffenseSource, DomainIPMapping, DomainName, ExternalData
from modules import AlchemyEncoder
from modules import ArielSearcher

# Resolver classes
from resolvers.Resolver import Resolver
from resolvers.Geo import Geo
from resolvers.Dns_whois import Dns_whois
from resolvers.SSL import SSL
#

logger = logging.getLogger('domainradar')

class DataResolver:
    def __init__(self, flask_app):
        """
        ! Constructor of the QLoader class
        @param DomainRadar-related objects from flask_app.config['DR']
        """
        self.flask_app = flask_app
        self.thread = None

        dr_objects = flask_app.config['DR']
        self.dr_config = dr_objects['CONFIG']

        self.resolving = False
        self.last_result = None
        self.current_result = {
            'success': True,
            'error_description': '',
            'domains_resolved_new': 0,
            'domains_resolved_total': 0,
            'domains_remaining': 0,
            'domains_total': 0
        }

        # Resolvers to use
        self.resolvers = []
        self.resolvers.append(Geo())
        self.resolvers.append(Dns_whois())
        self.resolvers.append(SSL())


    def startResolving(self):
        """
        ! Starts the loading process (controlled by Flask App)
        """
        self.resolving = True


    def isResolving(self):
        """
        ! Checks if there is a data resolution in progress (controlled by Flask App)
        @return True if the data resolution is in progress, False otherwise
        """
        return self.resolving


    def getLastResult(self):
        """
        ! Returns the result of the last run (controlled by Flask App)
        """
        return self.last_result


    def mainLoop(self):
        """
        ! Main loop of the DataResolver (controlled by a separate thread)
        """
        while True:
            if self.resolving == True:
                logger.info("DataResolver: Started resolution of external data...")

                domains_resolved_new = 0
                domains_resolved_total = 0
                domains_remaining = 0
                domains_total = 0

                # Get the list of domains that require addition data
                with self.flask_app.app_context():
                    domains = [r.domain_name for r in
                        DomainName.query.filter(DomainName.analyzed == false()).filter(DomainName.resolved == false()).all()]

                # Calculate how many domains need resolution
                domains_to_resolve = len(domains)

                logger.info(str(len(domains)) + " total unanalyzed domains with unresolved data.")

                # Prepare the thread executor
                executor = concurrent.futures.ThreadPoolExecutor()

                # Initialize dictionary for returned data
                domain_data = dict()
                for d in domains:
                    domain_data[d] = dict()

                # Run resolving threads
                for i in range(len(domains)):
                    executor.submit(self.__resolve_domain, domains[i], domain_data)

                # Wait for all threads to finish
                executor.shutdown(wait=True, cancel_futures=False)
                logger.info("DataResolver: Resolution done, verifying and saving results.")

                # Check results for each domain and save them to db
                for domain in domain_data:
                    resolution_ok = True
                    for resolver in domain_data[domain]:
                        result = domain_data[domain][resolver]

                        if result['success'] == False:
                            resolution_ok = False
                            resolver_name = result['resolver_name']
                            err = result['error_description']
                            logger.error("DataResolver: " + str(resolver_name) + ": " + str(err))
                        else:
                            with self.flask_app.app_context():
                                data_record = ExternalData(domain_name = domain, data_type = str(resolver), contents = json.dumps(result, default=str))
                                try:
                                    db.session.add(data_record)
                                    db.session.commit()
                                except exc.SQLAlchemyError as e:
                                    logger.error(e)
                                    db.session.rollback()

                    if not resolution_ok:
                        # Continue to new domain
                        continue

                    with self.flask_app.app_context():
                        domain_name_record = DomainName.query.filter(DomainName.domain_name == domain).one()
                        domain_name_record.resolved = True
                        try:
                            db.session.add(data_record)
                            db.session.commit()
                        except exc.SQLAlchemyError as e:
                            logger.error(e)
                            db.session.rollback()
                    domains_resolved_new += 1

                # Success
                with self.flask_app.app_context():
                    domains_total = db.session.query(DomainName.id).count()
                    domains_remaining = domains_to_resolve - domains_resolved_new
                    domains_resolved_total = domains_total - domains_remaining

                self.current_result['domains_resolved_new'] = domains_resolved_new
                self.current_result['domains_resolved_total'] = domains_resolved_total
                self.current_result['domains_remaining'] = domains_remaining
                self.current_result['domains_total'] = domains_total
                self.last_result = self.current_result
                self.resolving = False

                logger.info("DataResolver: Resolved and saved data for " + str(domains_resolved_new) + " domains in this run.")
                logger.info("DataResolver: Resolved domains total: " + str(domains_resolved_total))
                logger.info("DataResolver: Domains remaining to resolve: " + str(domains_remaining))
                logger.info("DataResolver: Domains total: " + str(domains_total))
                logger.info("DataResolver: Finished resolution of domain external data.")

            # Wait a second for the next iteration
            sleep(1)

    def __resolve_domain(self, domain_name, domain_data):


        # Find all related IP addresses for the domain
        with self.flask_app.app_context():
            ips = [r.ip_address for r in
                DomainIPMapping.query.filter(DomainIPMapping.domain_name == domain_name).all()]

        # Run resolvers for all external data types
        for resolver in self.resolvers:
            resolver_name = resolver.getName()

            # Resolve the external data
            data = resolver.resolve(domain_name, ips)

            # Save data for the given domain and resolver
            # NOTE: Deepcopy is needed here, otherwise next runs will change the referenced object (data)
            domain_data[domain_name][resolver_name] = copy.deepcopy(data)
