from services.vulnerability_intelligence.handlers.base_handler import BaseHandler
from models.vulnerability_intelligence import VulnerabilityIntelligence
from dateutil import parser as dateutil_parser

class CisaKevHandler(BaseHandler):
    def apply(self, vuln_intelligence: VulnerabilityIntelligence):
        try:
            cisa_notes = self.data.get('cisa_notes', '')
            if cisa_notes:
                vuln_intelligence.reference_urls.update(cisa_notes.split(' ; '))

            cisa_description = self.data.get('cisa_description')
            if cisa_description:
                formatted_description = cisa_description.replace('\n', ' ')
                if not any(desc["text"] == formatted_description for desc in vuln_intelligence.descriptions):
                    vuln_intelligence.descriptions.append({
                        "source": self.enrich_source_name("CISA KEV"),
                        "text": formatted_description,
                        "date": self._parse_date(self.data.get('cisa_dateAdded'))
                    })

            cisa_vendor = self.data.get('cisa_vendorProject')
            if cisa_vendor:
                vuln_intelligence.tags.add(cisa_vendor)
            cisa_product = self.data.get('cisa_product')
            if cisa_product:
                vuln_intelligence.vulnerable_components.add(cisa_product)

            cisa_cwes = self.data.get('cisa_cwes', [])
            if cisa_cwes:
                if isinstance(cisa_cwes, list):
                    vuln_intelligence.weaknesses.update(cisa_cwes)
                else:
                    vuln_intelligence.weaknesses.add(cisa_cwes)

            cisa_required_action = self.data.get('cisa_requiredAction')
            if cisa_required_action:
                formatted_mitigation = cisa_required_action.replace('\n', ' ')
                vuln_intelligence.descriptions.append({
                    "source": self.enrich_source_name_mitigation("CISA KEV"),
                    "text": formatted_mitigation,
                    "date": self._parse_date(self.data.get('cisa_dueDate'))
                })

        except Exception as e:
            print(f"[!] Error applying CISA KEV enrichment: {e}")

    def _parse_date(self, date_str: str) -> str:
        if not date_str:
            return "N/A"
        try:
            parsed_date = dateutil_parser.parse(date_str)
            return parsed_date.strftime("%Y-%m-%d")
        except Exception as e:
            print(f"[!] Error parsing date '{date_str}': {e}")
            return "N/A"