import elasticsearch
import urllib3
from elasticsearch_dsl import Search
import os
from forward_evidence.resource_id_mapper import map_resource_id

WAZUH_CHECK_INTERVAL = os.environ.get("wazuh_check_interval")
WAZUH_RULE_LEVEL = int(os.environ.get("wazuh_rule_level"))

class Checker:
    def __init__(self, wc, es, logger):
        self.wc = wc
        self.es = es
        self.logger = logger
        
    # Check if syscheck enabled
    def check_syscheck(self, agent):
        body = self.wc.req('GET', 'agents/' + agent[0] + '/config/syscheck/syscheck')
        
        measurement_result = body['data']['syscheck']['disabled'] == 'no'
        
        return body, measurement_result
    
    # Check if rootcheck enabled
    def check_rootcheck(self, agent):
        body = self.wc.req('GET', 'agents/' + agent[0] + '/config/syscheck/rootcheck')
        
        measurement_result = body['data']['rootcheck']['disabled'] == 'no'
        
        return body, measurement_result
    
    # Check if there's at least one valid alerting service
    def check_alert_integrations(self):
        body = self.wc.req('GET', 'manager/configuration')
        
        # Check email notifications integration
        try:
            email_notifications = (
                True if body['data']['affected_items'][0]['global']['email_notification'] == 'yes' else False)
        except:
            email_notifications = False
        
        # Check Slack and PagerDuty notifications integration
        try:
            integrations = body['data']['affected_items'][0]['integration']
            
            slack_notifications = pagerduty_notifications = False
            
            for integration in integrations:
                if integration['name'] == 'slack':
                    slack_notifications = True
                
                if integration['name'] == 'pagerduty':
                    pagerduty_notifications = True
        except:
            slack_notifications = pagerduty_notifications = False
        
        measurement_result = email_notifications or slack_notifications or pagerduty_notifications
        
        return body, measurement_result
    
    # Check for VirusTotal integration
    def check_virus_total_integration(self):
        body = self.wc.req('GET', 'manager/configuration')
        
        # Check VirusTotal integration
        try:
            integrations = body['data']['affected_items'][0]['integration']
            
            measurement_result = False
            
            for integration in integrations:
                if integration['name'] == 'virustotal':
                    measurement_result = True
                    break
        except:
            measurement_result = False
        
        return body, measurement_result
    
    # Check if ClamAV daemon process running
    def check_clamd_process(self, agent):
        body = self.wc.req('GET', 'syscollector/' + agent[0] + '/processes')
        
        measurement_result = False
        
        for package in body['data']['affected_items']:
            if package['name'] == 'clamd':
                measurement_result = True
                break
        
        return body, measurement_result
    
    # Check ClamAV logs in Elasticsearch
    def check_clamd_logs_elastic(self, agent):
        s = Search(using=self.es, index="wazuh-alerts-*") \
            .query("match", predecoder__program_name="clamd") \
            .query("match", rule__description="Clamd restarted") \
            .query("match", agent__id=agent[0])
        
        try:
            body = s.execute().to_dict()
        except (elasticsearch.exceptions.ConnectionError, TimeoutError, urllib3.exceptions.NewConnectionError, 
        		urllib3.exceptions.MaxRetryError) as err:
            self.logger.error(err)
            self.logger.error("Elasticsearch not available")

            return None, False

        measurement_result = len(body['hits']['hits']) > 0
        
        return body, measurement_result


    def check_security_events(self, agent):
        query = {
            "query": {
                "bool": {
                    "must": [
                        {
                            "match": {
                                "agent.id": agent[0]
                            }
                        },
                        {
                            "range" : {
                                "rule.level" : {
                                    "gte" : WAZUH_RULE_LEVEL
                                }
                            }
                        },
                        {
                            "range" : {
                                "@timestamp" : { 
                                    "gte" : "now-" + WAZUH_CHECK_INTERVAL + "s"
                                }
                            }
                        }
                    ]
                }
            }
        }
        
        try:
            body = self.es.search(index="wazuh-alerts-*", body=query)
        except (elasticsearch.exceptions.ConnectionError, TimeoutError, urllib3.exceptions.NewConnectionError, 
        		urllib3.exceptions.MaxRetryError) as err:
            self.logger.error(err)
            self.logger.error("Elasticsearch not available")

            return None

        self.logger.debug(map_resource_id(agent[1]) + " security events count: "  + str(len(body['hits']['hits'])))

        return len(body['hits']['hits'])