Spaces:
Sleeping
Sleeping
| 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 | |