import logging
import requests

LOGGER = logging.getLogger("mcp-cve")

class CVE:
    def __init__(
            self, 
            cve_id = None,
            token = None,
        ):
        self.cve_id = cve_id
        self.cve_url = "https://cveawg.mitre.org/api/cve/"
        self.poc_url = "https://raw.githubusercontent.com/adysec/nuclei_poc/refs/heads/main/poc/cve/%s.yaml"
        self.token = token
        LOGGER.debug(f"Initializing CVE with CVE ID: {cve_id}")

    def get_cve(self) -> dict:
        LOGGER.info(f"Fetching CVE: {self.cve_id}...")

        try:
            response = requests.get(self.cve_url + self.cve_id)

            if response.status_code == 200:
                LOGGER.info(f"Response: {response.status_code}")
                data = response.json()
                data["nuclei_poc"] = self.get_poc()
                return data
            else:
                LOGGER.info(f"CVE {self.cve_id} not found in NVD database.")
                return None

        except Exception as e:
            LOGGER.info(f"Error: {e}")
            return None

    def get_poc(self) -> str:
        LOGGER.info(f"Fetching POC: {self.cve_id}...")

        try:
            response = requests.get(self.poc_url % self.cve_id)

            if response.status_code == 200:
                LOGGER.info(f"Response: {response.status_code}")
                data = response.text
                return data
            else:
                LOGGER.info(f"POC {self.cve_id} not found in NVD database.")
                return None

        except Exception as e:
            LOGGER.info(f"Error: {e}")
            return None