Blame for app/utils.py

import os
import logging
from datetime import datetime, timedelta
import aiofiles
import json
import string
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"
PROMPT_DIR = "app/prompts"
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.")
            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 _prompt_base_path() -> str:
    return PROMPT_DIR

def _prompt_path(prompt_name: str, suffix=".md") -> str:
    return os.path.join(PROMPT_DIR, f"{prompt_name}{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 generate_openai_report(request_object: dict, format: str = 'json'):
    logger.info(f"Generating report in {format} format")
    response_format = { "type": "text" }
    system_prompt: str = ''
    if request_object['prompt']['source'] == 'name':
        template = get_prompt_template(_prompt_path(request_object['prompt']['name']))
        if format == 'json':
            name, ext = os.path.splitext(_prompt_path(request_object['prompt']['name']))
            system_prompt = get_prompt_template(f"{name}_json{ext}")
            response_format = { "type": "json_object" }

    else:
        template = request_object['prompt']['prompt']
    prompt = safe_format(template, request_object['prompt']['params'])
    response = await client.chat.completions.create(
        model = "gpt-4o-search-preview",
        web_search_options={
            "search_context_size": "low",
        },
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ],
#        response_format=response_format, # Not supported with web search
    )
    return response.choices[0].message.content

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}")
    
def safe_format(template: str, params: dict) -> str:
    # Filling a prompt string with parameters
    logger.info(f"Filling prompt string: {template} with parameters: {params}")
    formatter = string.Formatter()
    try:
        variables = [field_name for _, field_name, _, _ in formatter.parse(template) if field_name]
    except ValueError:
        raise ValueError(f'Invalid template in prompt_prompt')
    if not variables:
        return template
    missing = [var for var in variables if var not in params]
    if missing:
            raise KeyError(f'Missing keys for template: {missing}')
    filled_template = template.format_map(params)
    logger.info(f"Filled prompt string with the parameters: {filled_template}")
    return filled_template