from sator_core.models.vulnerability.attributes import VulnerabilityAttributes
from sator_core.models.vulnerability.description import VulnerabilityDescription

from sator_core.ports.driven.persistence.storage import StoragePersistencePort
from sator_core.ports.driven.extraction.attributes.vulnerability import VulnerabilityAttributesExtractorPort
from sator_core.ports.driving.extraction.attributes.vulnerability import VulnerabilityAttributesExtractionPort


class VulnerabilityAttributesExtraction(VulnerabilityAttributesExtractionPort):
    def __init__(self, storage_port: StoragePersistencePort, attributes_extractor: VulnerabilityAttributesExtractorPort):
        self.storage_port = storage_port
        self.attributes_extractor = attributes_extractor

    def extract_vulnerability_attributes(self, vulnerability_id: str) -> VulnerabilityAttributes | None:
        attributes = self.storage_port.load(VulnerabilityAttributes, vulnerability_id)

        if attributes:
            return attributes

        description = self.storage_port.load(VulnerabilityDescription, vulnerability_id)

        if description:
            attributes = self.attributes_extractor.extract_vulnerability_attributes(description)

            if attributes:
                attributes = VulnerabilityAttributes(**attributes)
                self.storage_port.save(attributes, vulnerability_id)
                return attributes

        return None
