import logging
import bisect
import re

from collections import defaultdict
from itertools import chain

from ..common.common import SectionHandler
from ..section_handlers.document_acknowledgments import Acknowledgments
from ..section_handlers.references import References
from ..section_handlers.notes import Notes
from ..common.utils import get_utc_timestamp


class Vulnerability(SectionHandler):
    # mapping of CVRF XML elements to CSAF json keys
    cvss_v3_mapping = {'BaseScoreV3': 'baseScore',
                       'TemporalScoreV3': 'temporalScore',
                       'EnvironmentalScoreV3': 'environmentalScore',
                       'VectorV3': 'vectorString'}

    cvss_v2_mapping = {'BaseScoreV2': 'baseScore',
                       'TemporalScoreV2': 'temporalScore',
                       'EnvironmentalScoreV2': 'environmentalScore',
                       'VectorV2': 'vectorString'}

    def __init__(self, config):
        super().__init__()
        self.csaf = []
        self.config = config
        self.remove_CVSS_values_without_vector = config['remove_CVSS_values_without_vector']
        self.default_CVSS3_version = config['default_CVSS3_version']

    def _process_mandatory_elements(self, root_element):
        pass

    @staticmethod
    def _handle_involvements(root_element):
        involvements = []
        for involvement_elem in root_element.Involvement:
            involvement = {'party': involvement_elem.attrib['Party'].lower(),
                           'status': involvement_elem.attrib['Status'].lower().replace(' ', '_')}

            if hasattr(involvement_elem, 'Description'):
                involvement['summary'] = involvement_elem.Description.text
            involvements.append(involvement)

        return involvements

    @staticmethod
    def _handle_product_statuses(root_element):
        statuses = defaultdict(list)
        for status_elem in root_element.Status:
            status_type = status_elem.attrib['Type'].lower().replace(' ', '_')
            product_ids = [product_id.text for product_id in status_elem.ProductID]
            statuses[status_type].extend(product_ids)

        return statuses

    @staticmethod
    def _handle_threats(root_element):
        threats = []
        for threat_elem in root_element.Threat:
            threat = {'details': threat_elem.Description.text,
                      'category': threat_elem.attrib['Type'].lower().replace(' ', '_')}

            if hasattr(threat_elem, 'ProductID'):
                threat['product_ids'] = [product_id.text for product_id in threat_elem.ProductID]

            if hasattr(threat_elem, 'GroupID'):
                threat['group_ids'] = [group_id.text for group_id in threat_elem.GroupID]

            if 'Date' in threat_elem.attrib:
                threat['date'] = get_utc_timestamp(threat_elem.attrib['Date'])

            threats.append(threat)

        return threats

    @staticmethod
    def _handle_remediations(root_element):

        category_mapping = {'Workaround': 'workaround',
                            'Mitigation': 'mitigation',
                            'Vendor Fix': 'vendor_fix',
                            'None Available': 'none_available',
                            'Will Not Fix': 'no_fix_planned'}

        remediations = []
        for remediation_elem in root_element.Remediation:
            remediation = {'category': category_mapping[remediation_elem.attrib['Type']],
                           'details': remediation_elem.Description.text}

            if hasattr(remediation_elem, 'Entitlement'):
                remediation['entitlements'] = [entitlement.text for entitlement in remediation_elem.Entitlement]

            if hasattr(remediation_elem, 'URL'):
                remediation['url'] = remediation_elem.URL.text

            if hasattr(remediation_elem, 'ProductID'):
                remediation['product_ids'] = [product_id.text for product_id in remediation_elem.ProductID]

            if hasattr(remediation_elem, 'GroupID'):
                remediation['group_ids'] = [group_id.text for group_id in remediation_elem.GroupID]

            if 'Date' in remediation_elem.attrib:
                remediation['date'] = get_utc_timestamp(remediation_elem.attrib['Date'])

            remediations.append(remediation)

        return remediations

    @staticmethod
    def _base_score_to_severity(base_score):
        base_severity = ((0, 'NONE'), (3.9, 'LOW'), (6.9, 'MEDIUM'), (8.9, 'HIGH'), (10, 'CRITICAL'))
        return base_severity[bisect.bisect_right(base_severity, (base_score, ''))][1]

    def _parse_score_set(self, score_set_element, mapping, version, json_property, product_status):
        """Parses ScoreSetV2 or ScoreSetV3 element."""

        # Parse all input elements except ProductID
        # note: baseScore is always present since it's mandatory for the input
        cvss_score = {csaf: getattr(score_set_element, cvrf).text for cvrf, csaf in mapping.items() if hasattr(score_set_element, cvrf)}

        # Convert all possible scores to float
        scores = ['baseScore', 'temporalScore', 'environmentalScore']
        for score in scores:
            if cvss_score.get(score):
                cvss_score[score] = float(cvss_score[score])

        # Only cvss_v3 has baseSeverity
        if json_property == 'cvss_v3':
            cvss_score['baseSeverity'] = self._base_score_to_severity(cvss_score['baseScore'])

        # HANDLE product_ids
        product_ids = []
        # if we have ProductID element(s), parse and use
        if hasattr(score_set_element, 'ProductID'):
            product_ids = [product_id.text for product_id in score_set_element.ProductID]
        # one of the conversion rules specifies what to do in case of missing ProductID element
        elif product_status:
            states = ('known_affected', 'first_affected', 'last_affected')
            product_ids = list(set(product_id for state in states for product_id in product_status.get(state, [])))

        if len(product_ids) == 0:
            SectionHandler.error_occurred = True
            logging.error('No product_id entry for CVSS score set.')

        # HANDLE vectorString
        # if missing, conversion fails unless remove_CVSS_values_without_vector is true
        # if remove_CVSS_values_without_vector is true, we just ignore the score_set
        if 'vectorString' not in cvss_score:
            if self.remove_CVSS_values_without_vector:
                logging.warning('No CVSS vector string found on the input, ignoring ScoreSet element due to'
                                '"remove_CVSS_values_without_vector" option.')
                return None
            else:
                SectionHandler.error_occurred = True
                logging.error('No CVSS vector string found on the input.')

        # DETERMINE CVSS v 3.x from namespace
        cvss_3_regex = '.*cvss-v(3\.[01]).*'
        match = re.match(cvss_3_regex, score_set_element.tag)
        if match:
            version = match.groups()[0]

        # DETERMINE CVSS v 3.x from vector if present
        if 'vectorString' in cvss_score and json_property == 'cvss_v3':
            # Regex for determining the CVSS version
            regex = "CVSS:(3\.[01]).*"
            match = re.match(regex, cvss_score['vectorString'])
            if not match:
                SectionHandler.error_occurred = True
                logging.error(f'CVSS vector {cvss_score["vectorString"]} is not valid.')
            else:
                version = match.groups()[0]

        cvss_score['version'] = version

        score = {json_property: cvss_score,
                 'products': product_ids}

        return score

    @staticmethod
    def _remove_cvssv3_duplicates(scores):
        """
        Removes products/cvssv3.x score sets for products having both v3.0 and v3.1 score. Three step approach:
         - find products having both versions specified
         - remove those products from score set with version 3.0
         - removes score sets with no products
        """
        # STEP 1
        products_v3_1 = set(chain.from_iterable([score_set['products'] for score_set in scores if
                                                 'cvss_v3' in score_set and score_set['cvss_v3']['version'] == '3.1']))
        products_v3_0 = set(chain.from_iterable([score_set['products'] for score_set in scores if
                                                 'cvss_v3' in score_set and score_set['cvss_v3']['version'] == '3.0']))

        both_versions = products_v3_0.intersection(products_v3_1)

        # STEP 2
        for score_set in scores:
            if 'cvss_v3' in score_set and score_set['cvss_v3']['version'] == '3.0':
                score_set['products'] = [product for product in score_set['products'] if product not in both_versions]

        # STEP 3
        scores = [score_set for score_set in scores if len(score_set['products']) > 0]

        return scores


    def _handle_scores(self, root_element, product_status):
        scores = []

        score_variants = (
            ('ScoreSetV2', self.cvss_v2_mapping, '2.0', 'cvss_v2'),
            ('ScoreSetV3', self.cvss_v3_mapping, self.default_CVSS3_version, 'cvss_v3'),
        )
        for score_variant, mapping, score_version, target in score_variants:
            for score_set in root_element.findall(f'{{*}}{score_variant}'):
                score = self._parse_score_set(score_set, mapping, score_version, target, product_status)
                if score is not None:
                    scores.append(score)

        scores = self._remove_cvssv3_duplicates(scores)

        return scores

    def _process_optional_elements(self, root_element):
        vulnerability = {}

        if hasattr(root_element, 'Acknowledgments'):
            # reuse Acknowledgments handler
            acknowledgments_handler = Acknowledgments()
            acknowledgments_handler.create_csaf(root_element.Acknowledgments)
            vulnerability['acknowledgments'] = acknowledgments_handler.csaf

        if hasattr(root_element, 'CVE'):
            # "^CVE-[0-9]{4}-[0-9]{4,}$" differs from the CVRF regex. Will be checked by json schema validation.
            vulnerability['cve'] = root_element.CVE.text

        if hasattr(root_element, 'CWE'):
            CWE_elements = root_element.CWE
            if len(CWE_elements) > 1:
                logging.warning(f'{len(CWE_elements)} CWE elements found, using only the first one.')
            vulnerability['cwe'] = {'id': CWE_elements[0].attrib['ID'],
                                    'name': CWE_elements[0].text}

        if hasattr(root_element, 'DiscoveryDate'):
            vulnerability['discovery_date'] = get_utc_timestamp(root_element.DiscoveryDate.text)

        if hasattr(root_element, 'ID'):
            vulnerability['id'] = {'system_name': root_element.ID.attrib['SystemName'],
                                   'text': root_element.ID.text}

        if hasattr(root_element, 'Involvements'):
            vulnerability['involvements'] = self._handle_involvements(root_element.Involvements)

        if hasattr(root_element, 'Notes'):
            # reuse Notes handler
            notes_handler = Notes()
            notes_handler.create_csaf(root_element.Notes)
            vulnerability['notes'] = notes_handler.csaf

        if hasattr(root_element, 'ProductStatuses'):
            vulnerability['product_status'] = self._handle_product_statuses(root_element.ProductStatuses)

        if hasattr(root_element, 'References'):
            # reuse References handler
            references_handler = References(config=self.config)
            references_handler.create_csaf(root_element.References)
            vulnerability['references'] = references_handler.csaf

        if hasattr(root_element, 'ReleaseDate'):
            vulnerability['release_date'] = get_utc_timestamp(root_element.ReleaseDate.text)

        if hasattr(root_element, 'Remediations'):
            vulnerability['remediations'] = self._handle_remediations(root_element.Remediations)

        if hasattr(root_element, 'CVSSScoreSets'):
            vulnerability['scores'] = self._handle_scores(root_element.CVSSScoreSets,
                                                          vulnerability.get('product_status'))

        if hasattr(root_element, 'Threats'):
            vulnerability['threats'] = self._handle_threats(root_element.Threats)

        if hasattr(root_element, 'Title'):
            vulnerability['title'] = root_element.Title.text

        self.csaf.append(vulnerability)
