Spaces:

dn-scribe
/

cas

Sleeping

cas / app /utils.py

add nvd-mitre fallback

1f07bde unverified over 1 year ago

9.82 kB

import os
import logging
from datetime import datetime, timedelta
import aiofiles
import json
from openai import AsyncOpenAI
import httpx

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# get working directory
current_dir = os.path.dirname(os.path.abspath(__file__))
if "/Users/" in current_dir:
    # If running on MacOS, set the report directory to a specific path
    REPORT_DIR = "tmp/reports"
else:
    REPORT_DIR = "/tmp/reports"
os.makedirs(REPORT_DIR, exist_ok=True)

client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Add a helper function to validate CVE-ID
def validate_cve_id(cve_id: str):
    parts = cve_id.upper().split("-")
    if len(parts) != 3 or not parts[2].isdigit():
        return False
    return True

def is_safe_path(base_path: str, target_path: str) -> bool:
    """
    Validate that the target path is within the base path and does not involve directory traversal.
    Args:
        base_path (str): The base directory path.
        target_path (str): The target file path to validate.
    Returns:
        bool: True if the target path is safe, False otherwise.
    """
    # Resolve absolute paths
    base_path = os.path.abspath(base_path)
    target_path = os.path.abspath(target_path)

    # Check if the target path starts with the base path
    return os.path.commonpath([base_path]) == os.path.commonpath([base_path, target_path])

async def get_cve_info(cve_id: str) -> dict | None:
    """
    Fetch CVE JSON data from cisagov/vulnrichment repo.
    Args:
        cve_id (str): The CVE identifier, e.g., "CVE-2025-1234".
    Returns:
        dict | None: Parsed CVE data or None if not found.
    """
    try:
        parts = cve_id.upper().split("-")
        if len(parts) != 3 or not parts[2].isdigit():
            logger.warning(f"Invalid CVE ID format: {cve_id}")
            return None

        year, cve_num = parts[1], parts[2].lstrip("0")
        padded = parts[2].zfill(4)
        num_len = len(padded)

        if num_len == 4:
            group = f"{padded[0]}xxx"
        elif num_len == 5:
            group = f"{padded[:2]}xxx"
        elif num_len == 6:
            group = f"{padded[:3]}xxx"
        else:
            logger.warning(f"Unsupported CVE numeric length: {cve_id}")
            return None

        url = f"https://raw.githubusercontent.com/cisagov/vulnrichment/develop/{year}/{group}/CVE-{year}-{padded}.json"
        logger.info(f"Fetching CVE info from {url}")
        async with httpx.AsyncClient(timeout=5.0) as client:
            resp = await client.get(url)
        logger.info(f"Response status code: {resp}")

        if resp.status_code == 200:
            return resp.json()
        elif resp.status_code == 404:
            logger.info(f"CVE {cve_id} not found in vulnrichment, falling back to MITRE")
            nvd_url = f"https://cveawg.mitre.org/api/cve/{cve_id}"
            async with httpx.AsyncClient(timeout=5.0) as client:
                resp = await client.get(nvd_url)
            if resp.status_code == 200:
                return resp.json()
            else:
                logger.error(f"Failed to fetch CVE from NVD: {resp.status_code} for {nvd_url}")
            return None
        else:
            logger.error(f"Unexpected status {resp.status_code} for {url}")
            return None

    except Exception as e:
        logger.exception(f"Failed to fetch CVE info for {cve_id}: {e}")
        return None

def _report_base_path() -> str:
    return REPORT_DIR

def _report_path(cve_id: str, suffix=".md") -> str:
    return os.path.join(REPORT_DIR, f"{cve_id}{suffix}")

def _lock_path(cve_id: str) -> str:
    return os.path.join(REPORT_DIR, f".{cve_id}.lock")

async def report_exists(cve_id: str, suffix: str = ".md") -> bool:
    path = _report_path(cve_id, suffix)
    if not os.path.exists(path):
        return False
    mtime = datetime.fromtimestamp(os.path.getmtime(path))
    return datetime.now() - mtime < timedelta(days=5)

async def load_report(cve_id: str, suffix: str = ".md") -> str:
    path = os.path.join(REPORT_DIR, f"{cve_id}{suffix}")
    logger.info(f"Loading report for cve_id={cve_id}, path={path}")
    async with aiofiles.open(path, "r") as f:
        content = await f.read()
        logger.info(f"Report loaded for cve_id={cve_id}, content length={len(content)}")
        return content if suffix == ".md" else json.loads(content)


async def write_report(cve_id: str, content, suffix: str = ".md"):
    path = os.path.join(REPORT_DIR, f"{cve_id}{suffix}")
    try:
        async with aiofiles.open(path, "w") as f:
            if suffix == ".md":
                await f.write(content)
            else:
                await f.write(json.dumps(content, indent=2))
        logger.info(f"✅ Report written: {path}")
    except Exception as e:
        logger.error(f"❌ Failed to write report to {path}: {e}")


async def get_cve_prompt(cve_id: str) -> str:
    template = get_prompt_template("app/prompts/cve_analysis.md")
    cve_info_text = ""
    cve_info = await get_cve_info(cve_id)
    if cve_info:
        cve_info_text = json.dumps(cve_info, indent=2)
    else:
        cve_info_text = f"Could not retrieve data for {cve_id}. Search the web for information: CVSS score, description, exploitability, affected versions and references"
    logging.info(f"Retrieved CVE info for {cve_id}: {cve_info_text}")
    # if not cve_info:
    #     raise ValueError(f"Could not retrieve data for {cve_id}")
    prompt = template.format(cve_id=cve_id, cve_info=cve_info_text)
    return prompt

def get_prompt_template(path: str) -> str:
    """
    Reads a prompt from a file and returns it as a string.
    Args:
        path (str): The file path to read the prompt from.
    Returns:
        str: The content of the file as a string.
    Raises:
        FileNotFoundError: If the file does not exist.
        IOError: If there is an error reading the file.
    """
    try:
        with open(path, "r", encoding="utf-8") as file:
            return file.read()
    except FileNotFoundError:
        logger.error(f"Prompt file not found: {path}")
        raise FileNotFoundError(f"The specified prompt file at '{path}' was not found. Please ensure the file exists and the path is correct.")
    except IOError as e:
        logger.error(f"Error reading prompt file {path}: {e}")
        raise IOError(f"An error occurred while reading the prompt file at '{path}': {e}")

async def generate_openai_stream(cve_id: str, json_result: bool = False):
    logger.info(f"Generating OpenAI stream for cve_id={cve_id}")
    editing_directions = "Provide the result as markdown text."
    if json_result:
        editing_directions = get_prompt_template("app/prompts/cve_analysis_json.md") # No args
        
    prompt = await get_cve_prompt(cve_id)
    stream = await client.chat.completions.create(
        model = "gpt-4o-search-preview",
        web_search_options={
            "search_context_size": "low",
        },
        messages=[
            {"role": "system", "content": editing_directions},
            {"role": "user", "content": prompt}],
        stream=True,
    )

    async for chunk in stream:
        delta = chunk.choices[0].delta.content or ""
        if delta:
            logger.debug(f"Stream chunk received for cve_id={cve_id}, chunk length={len(delta)}")
            yield delta

async def generate_openai_json(cve_id: str):
    logger.info(f"Generating OpenAI JSON for cve_id={cve_id}")
    prompt = await get_cve_prompt(cve_id)
    system_prompt = get_prompt_template("app/prompts/cve_analysis_json.md")
    model = "gpt-4o-search-preview"
    response = await client.chat.completions.create(
        model=model,
        web_search_options={
            "search_context_size": "low",
        },
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ],
    )
    text = response.choices[0].message.content

    if not text.strip():
        logger.error(f"OpenAI returned an empty response for cve_id={cve_id}")
        raise ValueError("OpenAI returned an empty response. Unable to generate JSON.")

    logger.info(f"OpenAI JSON generated for cve_id={cve_id}, response length={len(text)}")
    return text

async def ghsa_to_cve(ghsa_id: str) -> str:
    """
    Convert a GHSA ID to a CVE ID using OpenAI's API.
    Args:
        ghsa_id (str): The GHSA identifier, e.g., "GHSA-1234-5678-9abc".
    Returns:
        str: The corresponding CVE ID or an error message.
    """
    logger.info(f"Converting GHSA ID {ghsa_id} to CVE ID")
    prompt = f"Convert the following GHSA ID to a CVE ID: {ghsa_id}. Return ONLY the CVE ID (or None if not found)."
    model = "gpt-4o-search-preview"
    response = await client.chat.completions.create(
        model=model,
        web_search_options={
            "search_context_size": "low",
        },
        messages=[
            {"role": "user", "content": prompt}
        ],
    )

    # Log input and output token usage
    logger.info(f"Input tokens: {response.usage.prompt_tokens}, Output tokens: {response.usage.completion_tokens}")

    cve_id = response.choices[0].message.content

    if not cve_id.strip():
        logger.error(f"OpenAI returned an empty response for cve_id={cve_id}")
        raise ValueError("OpenAI returned an empty response. Unable to generate JSON.")
    
    cve_id = cve_id.rstrip()  # Remove any leading/trailing whitespace
    if validate_cve_id(cve_id):  # Validate CVE-ID
        logger.info(f"OpenAI converted GHSA  {ghsa_id} to CVE ID {cve_id}")
        return cve_id
    else:
        logger.error(f"OpenAI returned an invalid CVE ID: {cve_id}")
        raise ValueError(f"OpenAI returned an invalid CVE ID: {cve_id}")

	import os
	import logging
	from datetime import datetime, timedelta
	import aiofiles
	import json
	from openai import AsyncOpenAI
	import httpx

	# Configure logging
	logging.basicConfig(level=logging.INFO)
	logger = logging.getLogger(__name__)

	# get working directory
	current_dir = os.path.dirname(os.path.abspath(__file__))
	if "/Users/" in current_dir:
	# If running on MacOS, set the report directory to a specific path
	REPORT_DIR = "tmp/reports"
	else:
	REPORT_DIR = "/tmp/reports"
	os.makedirs(REPORT_DIR, exist_ok=True)

	client = AsyncOpenAI(api_key=os.getenv("OPENAI_API_KEY"))

	# Add a helper function to validate CVE-ID
	def validate_cve_id(cve_id: str):
	parts = cve_id.upper().split("-")
	if len(parts) != 3 or not parts[2].isdigit():
	return False
	return True

	def is_safe_path(base_path: str, target_path: str) -> bool:
	"""
	Validate that the target path is within the base path and does not involve directory traversal.

	Args:
	base_path (str): The base directory path.
	target_path (str): The target file path to validate.

	Returns:
	bool: True if the target path is safe, False otherwise.
	"""
	# Resolve absolute paths
	base_path = os.path.abspath(base_path)
	target_path = os.path.abspath(target_path)

	# Check if the target path starts with the base path
	return os.path.commonpath([base_path]) == os.path.commonpath([base_path, target_path])

	async def get_cve_info(cve_id: str) -> dict \| None:
	"""
	Fetch CVE JSON data from cisagov/vulnrichment repo.

	Args:
	cve_id (str): The CVE identifier, e.g., "CVE-2025-1234".

	Returns:
	dict \| None: Parsed CVE data or None if not found.
	"""
	try:
	parts = cve_id.upper().split("-")
	if len(parts) != 3 or not parts[2].isdigit():
	logger.warning(f"Invalid CVE ID format: {cve_id}")
	return None

	year, cve_num = parts[1], parts[2].lstrip("0")
	padded = parts[2].zfill(4)
	num_len = len(padded)

	if num_len == 4:
	group = f"{padded[0]}xxx"
	elif num_len == 5:
	group = f"{padded[:2]}xxx"
	elif num_len == 6:
	group = f"{padded[:3]}xxx"
	else:
	logger.warning(f"Unsupported CVE numeric length: {cve_id}")
	return None

	url = f"https://raw.githubusercontent.com/cisagov/vulnrichment/develop/{year}/{group}/CVE-{year}-{padded}.json"
	logger.info(f"Fetching CVE info from {url}")
	async with httpx.AsyncClient(timeout=5.0) as client:
	resp = await client.get(url)
	logger.info(f"Response status code: {resp}")

	if resp.status_code == 200:
	return resp.json()
	elif resp.status_code == 404:
	logger.info(f"CVE {cve_id} not found in vulnrichment, falling back to MITRE")
	nvd_url = f"https://cveawg.mitre.org/api/cve/{cve_id}"
	async with httpx.AsyncClient(timeout=5.0) as client:
	resp = await client.get(nvd_url)
	if resp.status_code == 200:
	return resp.json()
	else:
	logger.error(f"Failed to fetch CVE from NVD: {resp.status_code} for {nvd_url}")
	return None
	else:
	logger.error(f"Unexpected status {resp.status_code} for {url}")
	return None

	except Exception as e:
	logger.exception(f"Failed to fetch CVE info for {cve_id}: {e}")
	return None

	def _report_base_path() -> str:
	return REPORT_DIR

	def _report_path(cve_id: str, suffix=".md") -> str:
	return os.path.join(REPORT_DIR, f"{cve_id}{suffix}")

	def _lock_path(cve_id: str) -> str:
	return os.path.join(REPORT_DIR, f".{cve_id}.lock")

	async def report_exists(cve_id: str, suffix: str = ".md") -> bool:
	path = _report_path(cve_id, suffix)
	if not os.path.exists(path):
	return False
	mtime = datetime.fromtimestamp(os.path.getmtime(path))
	return datetime.now() - mtime < timedelta(days=5)

	async def load_report(cve_id: str, suffix: str = ".md") -> str:
	path = os.path.join(REPORT_DIR, f"{cve_id}{suffix}")
	logger.info(f"Loading report for cve_id={cve_id}, path={path}")
	async with aiofiles.open(path, "r") as f:
	content = await f.read()
	logger.info(f"Report loaded for cve_id={cve_id}, content length={len(content)}")
	return content if suffix == ".md" else json.loads(content)


	async def write_report(cve_id: str, content, suffix: str = ".md"):
	path = os.path.join(REPORT_DIR, f"{cve_id}{suffix}")
	try:
	async with aiofiles.open(path, "w") as f:
	if suffix == ".md":
	await f.write(content)
	else:
	await f.write(json.dumps(content, indent=2))
	logger.info(f"✅ Report written: {path}")
	except Exception as e:
	logger.error(f"❌ Failed to write report to {path}: {e}")


	async def get_cve_prompt(cve_id: str) -> str:
	template = get_prompt_template("app/prompts/cve_analysis.md")
	cve_info_text = ""
	cve_info = await get_cve_info(cve_id)
	if cve_info:
	cve_info_text = json.dumps(cve_info, indent=2)
	else:
	cve_info_text = f"Could not retrieve data for {cve_id}. Search the web for information: CVSS score, description, exploitability, affected versions and references"
	logging.info(f"Retrieved CVE info for {cve_id}: {cve_info_text}")
	# if not cve_info:
	# raise ValueError(f"Could not retrieve data for {cve_id}")
	prompt = template.format(cve_id=cve_id, cve_info=cve_info_text)
	return prompt

	def get_prompt_template(path: str) -> str:
	"""
	Reads a prompt from a file and returns it as a string.

	Args:
	path (str): The file path to read the prompt from.

	Returns:
	str: The content of the file as a string.

	Raises:
	FileNotFoundError: If the file does not exist.
	IOError: If there is an error reading the file.
	"""
	try:
	with open(path, "r", encoding="utf-8") as file:
	return file.read()
	except FileNotFoundError:
	logger.error(f"Prompt file not found: {path}")
	raise FileNotFoundError(f"The specified prompt file at '{path}' was not found. Please ensure the file exists and the path is correct.")
	except IOError as e:
	logger.error(f"Error reading prompt file {path}: {e}")
	raise IOError(f"An error occurred while reading the prompt file at '{path}': {e}")

	async def generate_openai_stream(cve_id: str, json_result: bool = False):
	logger.info(f"Generating OpenAI stream for cve_id={cve_id}")
	editing_directions = "Provide the result as markdown text."
	if json_result:
	editing_directions = get_prompt_template("app/prompts/cve_analysis_json.md") # No args

	prompt = await get_cve_prompt(cve_id)
	stream = await client.chat.completions.create(
	model = "gpt-4o-search-preview",
	web_search_options={
	"search_context_size": "low",
	},
	messages=[
	{"role": "system", "content": editing_directions},
	{"role": "user", "content": prompt}],
	stream=True,
	)

	async for chunk in stream:
	delta = chunk.choices[0].delta.content or ""
	if delta:
	logger.debug(f"Stream chunk received for cve_id={cve_id}, chunk length={len(delta)}")
	yield delta

	async def generate_openai_json(cve_id: str):
	logger.info(f"Generating OpenAI JSON for cve_id={cve_id}")
	prompt = await get_cve_prompt(cve_id)
	system_prompt = get_prompt_template("app/prompts/cve_analysis_json.md")
	model = "gpt-4o-search-preview"
	response = await client.chat.completions.create(
	model=model,
	web_search_options={
	"search_context_size": "low",
	},
	messages=[
	{"role": "system", "content": system_prompt},
	{"role": "user", "content": prompt}
	],
	)
	text = response.choices[0].message.content

	if not text.strip():
	logger.error(f"OpenAI returned an empty response for cve_id={cve_id}")
	raise ValueError("OpenAI returned an empty response. Unable to generate JSON.")

	logger.info(f"OpenAI JSON generated for cve_id={cve_id}, response length={len(text)}")
	return text

	async def ghsa_to_cve(ghsa_id: str) -> str:
	"""
	Convert a GHSA ID to a CVE ID using OpenAI's API.

	Args:
	ghsa_id (str): The GHSA identifier, e.g., "GHSA-1234-5678-9abc".

	Returns:
	str: The corresponding CVE ID or an error message.
	"""
	logger.info(f"Converting GHSA ID {ghsa_id} to CVE ID")
	prompt = f"Convert the following GHSA ID to a CVE ID: {ghsa_id}. Return ONLY the CVE ID (or None if not found)."
	model = "gpt-4o-search-preview"
	response = await client.chat.completions.create(
	model=model,
	web_search_options={
	"search_context_size": "low",
	},
	messages=[
	{"role": "user", "content": prompt}
	],
	)

	# Log input and output token usage
	logger.info(f"Input tokens: {response.usage.prompt_tokens}, Output tokens: {response.usage.completion_tokens}")

	cve_id = response.choices[0].message.content

	if not cve_id.strip():
	logger.error(f"OpenAI returned an empty response for cve_id={cve_id}")
	raise ValueError("OpenAI returned an empty response. Unable to generate JSON.")

	cve_id = cve_id.rstrip() # Remove any leading/trailing whitespace
	if validate_cve_id(cve_id): # Validate CVE-ID
	logger.info(f"OpenAI converted GHSA {ghsa_id} to CVE ID {cve_id}")
	return cve_id
	else:
	logger.error(f"OpenAI returned an invalid CVE ID: {cve_id}")
	raise ValueError(f"OpenAI returned an invalid CVE ID: {cve_id}")