#!/usr/bin/env python3

"""QLoader.py: QRadar data loading daemon"""
__author__      = "Radek Hranicky"

import logging
import json
import time

from datetime import datetime
from time import sleep
from sqlalchemy import exc

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

logger = logging.getLogger('domainradar')

class QLoader:
    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.dr_qradar_api_client = dr_objects['QRADAR_API_CLIENT']
        self.dr_ariel_searcher = dr_objects['ARIEL_SEARCHER']

        self.loading = False
        self.last_result = None
        self.current_result = {
            'success': True,
            'error_description': '',
            'offense_sources_new': 0,
            'offense_sources_updated': 0,
            'offense_sources_total': 0,
            'domain_ip_mappings_new': 0,
            'domain_ip_mappings_total': 0,
            'unique_domains_new': 0,
            'unique_domains_total': 0
        }


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


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


    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 QLoader (controlled by a separate thread)
        """
        while True:
            if self.loading == True:
                logger.info("QLoader: Started loading data from QRadar SIEM...")

                # Loading offense sources
                logger.info("QLoader: STEP 1/2 - Offense sources")
                if self.loadOffenseSources():
                    logger.info("QLoader: Offense sources successfully loaded.")
                    self.current_result['success'] = True
                    self.current_result['error_description'] = ""
                else:
                    logger.error("QLoader: Offense sources loading failed!")
                    self.current_result['success'] = False
                    self.current_result['error_description'] = "Loading Offense sources failed. See log for details."
                    self.loading = False

                # Loading domain-ip mappings
                logger.info("QLoader: STEP 2/2 - Domain-IP mappings")
                if self.loadDomainIPMappings():
                    logger.info("QLoader: Domain-IP mappings successfuly loaded.")
                    self.current_result['success'] = True
                    self.current_result['error_description'] = ""
                else:
                    logger.error("QLoader: Loading Domain-IP mappings failed!")
                    self.current_result['success'] = False
                    self.current_result['error_description'] = "Loading Domain-IP mappings failed. See log for details."
                    self.loading = False
                    continue

                # Success
                self.last_result = self.current_result
                self.loading = False

                # Loading Domain-IP mappings
                logger.info("QLoader: Finished loading data from QRadar SIEM.")

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


    def loadOffenseSources(self):
        """
        ! Load offense IPs and metrics
        """
        logger.info("QLoader: Loading offense source IP addresses...")
        offense_sources = self.dr_qradar_api_client.request_get("/siem/source_addresses")

        if offense_sources == None:
            logger.error("QLoader: Unable to use QRadar API - check API token! ")
            return False

        logger.info("QLoader: Processing " + str(len(offense_sources)) + " offense sources from QRadar SIEM.")

        offense_sources_new = 0
        offense_sources_updated = 0
        offense_sources_total = 0

        with self.flask_app.app_context():
            for offense_source in offense_sources:
                source_ip = offense_source['source_ip']
                offense_count = len(offense_source['offense_ids'])
                event_flow_count = offense_source['event_flow_count']
                magnitude = offense_source['magnitude']

                # Check whether the source exists in the database
                existing_source = OffenseSource.query.filter(OffenseSource.source_ip == source_ip).first()
                if existing_source:
                    # Offense source exists, check if it was not changed
                    change = False
                    if existing_source.offense_count != offense_count:
                        existing_source.offense_count = offense_count
                        change = True
                    if existing_source.event_flow_count != event_flow_count:
                        existing_source.event_flow_count = event_flow_count
                        change = True
                    if existing_source.magnitude != magnitude:
                        existing_source.magnitude = magnitude
                        change = True
                    if change:
                    # Update the record
                        try:
                            db.session.commit()
                        except exc.SQLAlchemyError as e:
                            logger.error("QLoader: Error updating offense source record: " + e)
                            db.session.rollback()
                            return False
                        offense_sources_updated += 1
                else:
                    # This is a new offense source - it needs to be added to the database
                    os_record = OffenseSource(source_ip=source_ip, event_flow_count=event_flow_count, magnitude=magnitude, offense_count=offense_count)
                    try:
                        db.session.add(os_record)
                        db.session.commit()
                    except exc.SQLAlchemyError as e:
                        logger.error(e)
                        db.session.rollback()
                        return False
                    else:
                        offense_sources_new += 1

            offense_sources_total = db.session.query(OffenseSource.id).count()
            self.current_result['offense_sources_new'] = offense_sources_new
            self.current_result['offense_sources_updated'] = offense_sources_updated
            self.current_result['offense_sources_total'] = offense_sources_total


        logger.info("QLoader: " + str(offense_sources_new) + " new offense sources, " +
            str(offense_sources_updated) + " updated, " + str(offense_sources_total) + " total."
        )
        return True


    def loadDomainIPMappings(self):
        """
        ! Load offense Domain-IP mappings and unique domain names
        """
        logger.info("QLoader: Loading domain-IP mappings from Ariel db...")

        sid = self.dr_ariel_searcher.make_search("SELECT \"DNS rrname\", \"DNS rdata\" FROM events WHERE event_type='dns' AND \"DNS rrtype\"='A' AND \"DNS rdata\" != NULL LAST 30 DAYS")

        while not self.dr_ariel_searcher.is_completed(sid):
            logger.info("QLoader: Waiting for Ariel query to finish...")
            time.sleep(10)
        mappings = self.dr_ariel_searcher.results(sid)["events"]
        logger.info(str(len(mappings)) + " domain-IP mappings found in QRadar Ariel db.")
        logger.info("QLoader: Processing...")

        unique_domains_new = 0
        unique_domains_total = 0
        domain_ip_mappings_new = 0
        domain_ip_mappings_total = 0

        with self.flask_app.app_context():
            for mapping in mappings:
                ip = mapping["DNS rdata"]
                domain = mapping["DNS rrname"]
                # Check if the domain was seen before
                existing_domain = DomainName.query.filter(DomainName.domain_name == domain).first()
                if not existing_domain:
                    # New yet-unseen domain
                    domain_record = DomainName(domain_name=domain)
                    try:
                        db.session.add(domain_record)
                        db.session.commit()
                    except exc.SQLAlchemyError as e:
                        logger.error(e)
                        db.session.rollback()
                        return False
                    unique_domains_new += 1

                # Check if the mapping exists or not
                existing_mapping = DomainIPMapping.query.filter(DomainIPMapping.domain_name == domain).filter(DomainIPMapping.ip_address == ip).first()
                if not existing_mapping:
                    # New mapping
                    mapping_record = DomainIPMapping(domain_name = domain, ip_address = ip)
                    try:
                        db.session.add(mapping_record)
                        db.session.commit()
                    except exc.SQLAlchemyError as e:
                        logger.error(e)
                        db.session.rollback()
                        return False
                    domain_ip_mappings_new += 1

            unique_domains_total = db.session.query(DomainName.id).count()
            domain_ip_mappings_total = db.session.query(DomainIPMapping.id).count()

            self.current_result['unique_domains_new'] = unique_domains_new
            self.current_result['unique_domains_total'] = unique_domains_total
            self.current_result['domain_ip_mappings_new'] = domain_ip_mappings_new
            self.current_result['domain_ip_mappings_total'] = domain_ip_mappings_total

        logger.info("QLoader: " + str(unique_domains_new) + " new domain names, " +
            str(unique_domains_total) + " total."
        )
        logger.info("QLoader: " + str(domain_ip_mappings_new) + " new Domain-IP mappings, " +
            str(domain_ip_mappings_total) + " total."
        )
        return True
