cve_cli.py

import os
import json
import nvdlib
from openai import OpenAI
import logging
import datetime
from datetime import datetime
import re
from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.markdown import Markdown

client = None

# --- Logging Setup ---
LOG_DIR = "logs"
if not os.path.exists(LOG_DIR):
    os.makedirs(LOG_DIR)

log_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')

# File Handler
log_file = os.path.join(LOG_DIR, 'cve_cli.log')
file_handler = logging.FileHandler(log_file)
file_handler.setFormatter(log_formatter)
file_handler.setLevel(logging.INFO)

# Console Handler
console_handler = logging.StreamHandler()
console_handler.setFormatter(log_formatter)
console_handler.setLevel(logging.WARNING)

# Get root logger
logger = logging.getLogger()
logger.setLevel(logging.INFO)
logger.addHandler(file_handler)
logger.addHandler(console_handler)
# --- End Logging Setup ---

def load_api_keys():
    logger.info("Loading API keys from environment variables.")
    openai_key = os.getenv("OPENAI_API_KEY")
    nvd_api_key = os.getenv("NVD_API_KEY")
    if not openai_key or not nvd_api_key:
        logger.error("Missing OPENAI_API_KEY or NVD_API_KEY environment variable. Exiting.")
        exit(1)
    global client
    client = OpenAI(api_key=openai_key)
    logger.info("API keys loaded successfully.")
    return nvd_api_key

def extract_entities(question: str) -> dict:
    logger.debug(f"Attempting to extract entities from query: '{question}'")
    # Use a regular string since the variable is appended separately
    prompt_template = """Extract structured information from the security question below. Identify:
    - Specific CVE IDs (e.g., 'CVE-2021-44228')
    - Software/hardware components (e.g., 'log4j', 'Windows Server 2019')
    - Versions if specified (e.g., 'log4j 2.15.0')
    - Severity keywords (map to 'LOW', 'MEDIUM', 'HIGH', 'CRITICAL', e.g., 'critical issues' -> 'CRITICAL')
    - Relevant date ranges (YYYY-MM-DD format, e.g., 'since last month', 'in 2022')
    - Network context/architecture (e.g., 'DMZ', 'internal network', 'public facing', 'firewall connection', 'outdated system')
    - Difficulty level (e.g., 'novice', 'easy', 'medium', 'hard', 'master')

    Return ONLY a valid JSON object with keys: 'cve_id', 'components', 'versions', 'severity', 'start_date', 'end_date', 'network_context', 'difficulty'.
    If difficulty is not specified, default to 'medium'.
    If other values aren't found, use null or an empty list/dict.

    Example Input: 'Any critical log4j vulns since 2023 on my internal web server?'
    Example Output:
    ```json
    {{
      "cve_id": null,
      "components": ["log4j", "internal web server"],
      "versions": {{}},
      "severity": "CRITICAL",
      "start_date": "2023-01-01",
      "end_date": null,
      "network_context": {{"location": "internal", "type": "web server"}},
      "difficulty": "medium" # Default if not specified
    }}
    ```
    Question: "{question}"
    """
    prompt = prompt_template.format(question=question)
    resp = client.chat.completions.create(
        model="gpt-4.1",
        messages=[
            {"role": "system", "content": "You are an assistant extracting structured data for NVD searches. Map severity keywords and date ranges accurately."},
            {"role": "user", "content": prompt}
        ],
        temperature=0,
    )
    content = resp.choices[0].message.content.strip()
    logger.debug(f"Raw LLM response for entity extraction: '{content}'")

    # Strip markdown fences if present
    if content.startswith("```json"):
        content = content[7:] # Remove ```json
    if content.endswith("```"):
        content = content[:-3] # Remove ```
    content = content.strip() # Clean up any extra whitespace

    try:
        entities = json.loads(content)
        logger.info(f"Extracted entities: {json.dumps(entities)}")
        return entities
    except json.JSONDecodeError:
        logger.error(f"Failed to parse JSON from LLM entity extraction response: {content}")
        return {}

def strip_markdown(text):
    """Removes common markdown formatting for cleaner CLI output."""
    if not text:
        return ""
    # Remove headers (###, ##, #)
    text = re.sub(r'^#+\s+', '', text, flags=re.MULTILINE)
    # Remove bold (**text**)
    text = re.sub(r'\*\*(.*?)\*\*', r'\1', text)
    # Remove list items (*, -, +) - handles potential leading spaces
    text = re.sub(r'^\s*[\*\-\+]\s+', '', text, flags=re.MULTILINE)
    # Remove code blocks (```...```)
    text = re.sub(r'```.*?```', '', text, flags=re.DOTALL)
    # Remove inline code ticks (`code`)
    text = re.sub(r'`(.*?)`', r'\1', text)
    # Remove extra newlines
    text = re.sub(r'\n{2,}', '\n\n', text).strip()
    return text

def _parse_date_range(date_range_str: str) -> tuple:
    """Parses the date range string into start and end datetime objects."""
    if not date_range_str:
        return None, None

    now = datetime.now()
    start_date, end_date = None, now # Default end date is now

    date_range_str = date_range_str.lower()

    try:
        if 'last week' in date_range_str:
            start_date = now - datetime.timedelta(weeks=1)
        elif 'last month' in date_range_str:
            # Approximate month as 30 days for simplicity
            start_date = now - datetime.timedelta(days=30)
        elif 'last year' in date_range_str:
            start_date = now - datetime.timedelta(days=365)
        elif 'since ' in date_range_str:
            # Attempt to parse year like 'since 2022'
            year_str = date_range_str.split('since ')[-1].strip()
            if year_str.isdigit() and len(year_str) == 4:
                start_date = datetime(int(year_str), 1, 1)
            else:
                 # Try parsing a specific date like 'since YYYY-MM-DD'
                 try:
                     start_date = datetime.strptime(year_str, '%Y-%m-%d')
                 except ValueError:
                    logger.warning(f"Could not parse date string: {year_str}")
        elif 'between ' in date_range_str and ' and ' in date_range_str:
             # Attempt to parse 'between YYYY-MM-DD and YYYY-MM-DD'
            parts = date_range_str.split('between ')[-1].split(' and ')
            if len(parts) == 2:
                try:
                    start_date = datetime.strptime(parts[0].strip(), '%Y-%m-%d')
                    end_date = datetime.strptime(parts[1].strip(), '%Y-%m-%d')
                except ValueError:
                    logger.warning(f"Could not parse date range: {date_range_str}")
            else:
                logger.warning(f"Could not parse date range format: {date_range_str}")

        else:
            # Attempt basic date parsing if it's just a date string
            try:
                start_date = datetime.strptime(date_range_str, '%Y-%m-%d')
                # If only one date is given, assume it's the start date?
                # Or maybe treat it as a single day? Let's assume start date for now.
            except ValueError:
                 logger.warning(f"Unrecognized date range format: {date_range_str}")

    except Exception as e:
        logger.error(f"Error parsing date range '{date_range_str}': {e}")
        return None, None

    # Format for NVD API (YYYY-MM-DDTHH:mm:ss.SSSZ)
    # NVD uses UTC time
    utc_offset_str = datetime.now().astimezone().strftime('%z') # Get local offset like +HHMM or -HHMM
    start_date_str = start_date.strftime('%Y-%m-%dT00:00:00.000') + utc_offset_str if start_date else None
    end_date_str = end_date.strftime('%Y-%m-%dT%H:%M:%S.999') + utc_offset_str if end_date else None
    # NVD API v2 seems to prefer simpler ISO format like YYYY-MM-DDTHH:mm:ss for pubStartDate/pubEndDate
    # Let's try that format.
    start_date_nvd = start_date.isoformat(timespec='seconds') if start_date else None
    end_date_nvd = end_date.isoformat(timespec='seconds') if end_date else None

    logger.debug(f"Parsed date range: Start={start_date_nvd}, End={end_date_nvd}")
    return start_date_nvd, end_date_nvd


def _search_nvd_by_cve_id(cve_id_entity, api_key: str) -> list:
    """Searches NVD by a specific CVE ID."""
    search_cve_id = None
    if isinstance(cve_id_entity, list) and cve_id_entity:
        search_cve_id = cve_id_entity[0] # Use the first CVE ID if it's a list
    elif isinstance(cve_id_entity, str):
        search_cve_id = cve_id_entity

    if search_cve_id:
        try:
            logger.info(f"Querying NVD for specific CVE ID: {search_cve_id}")
            # Use the cleaned search_cve_id string
            results = nvdlib.searchCVE(cveId=search_cve_id, key=api_key, delay=1)
            logger.debug(f"NVD results for {search_cve_id}: {results}")
            return results
        except Exception as e:
            logger.error(f"NVDlib error searching for CVE ID '{search_cve_id}': {e}")
            return []
    else:
        logger.warning("Extracted CVE ID was present but invalid.")
        return []


def _search_nvd_by_keywords(entities: dict, api_key: str, difficulty: str) -> list:
    """Searches NVD by keywords, components, severity, and date range."""
    component = entities.get('component') # Should be a list now
    version = entities.get('version')
    severity_keywords = entities.get('severity') # This isn't directly used by NVD API v2 filters
    date_range_str = entities.get('date_range')

    # Map difficulty to CVSS V3 severity for NVD API filtering
    min_severity, max_severity = map_difficulty_to_cvss(difficulty)
    nvd_severity = map_difficulty_to_nvd_severity(difficulty)

    # Parse date range
    start_date, end_date = _parse_date_range(date_range_str)

    # Build search parameters
    search_params = {
        'delay': 1,
        'key': api_key,
        #'cvssV3Severity': nvd_severity, # Use range instead
        'cvssV3Metrics': f"V3::{min_severity}", # Filter by base score range potentially
        'resultsPerPage': 20 # Get a pool of results
    }
    # NVD API v2 doesn't have a direct component filter like v1.
    # We use keyword search and filter client-side or rely on CVE description matching.
    # Using 'keywordSearch' for component/version
    keyword_parts = []
    if component:
        # If component is a list, join elements for keyword search
        if isinstance(component, list):
            keyword_parts.extend(component)
        elif isinstance(component, str):
            keyword_parts.append(component)

    if version:
        keyword_parts.append(version)

    if keyword_parts:
        search_params['keywordSearch'] = " ".join(keyword_parts)
        search_params['keywordExactMatch'] = False # Allow partial matches
    else:
         # Avoid making a broad query without keywords
        logger.warning("No keywords (component/version) extracted for search.")
        return []

    # Add date filters if parsed successfully
    if start_date:
        search_params['pubStartDate'] = start_date
    if end_date:
        search_params['pubEndDate'] = end_date

    # Add CVSS v3 score range if mapped
    if min_severity:
         # This requires the API endpoint support - check NVD docs
         # Example: cvssV3Metric=BASE_SCORE_RANGE:[7.0, 10.0]
         # nvdlib might not directly support this; consider raw request if needed
         # For now, let's rely on the severity enum if available or post-filter
         search_params['cvssV3Severity'] = nvd_severity # Use the enum filter

    logger.info(f"Searching NVD with keywords: '{search_params.get('keywordSearch', 'N/A')}', severity: {nvd_severity}, dates: {start_date} to {end_date}")
    all_results = []
    try:
        results = nvdlib.searchCVE(**search_params)
        all_results.extend(results)
        logger.debug(f"Found {len(results)} results from NVD keyword search.")
    except Exception as e:
        logger.error(f"NVDlib error during keyword search: {e}")

    # Optional: Post-filtering based on severity keywords if needed
    # if severity_keywords and all_results:
    #     pass # Implement filtering logic if basic NVD filter isn't sufficient

    return all_results


# Simplified main query function
def query_nvd(entities: dict, api_key: str, difficulty: str = 'medium') -> list:
    """Queries the NVD based on extracted entities, delegating to helpers."""
    logger.debug(f"Querying NVD with entities: {entities}")
    cve_id_entity = entities.get('cve_id')

    # --- Prioritize CVE ID Search ---
    if cve_id_entity:
        results = _search_nvd_by_cve_id(cve_id_entity, api_key)
        if results: # Return immediately if CVE ID search is successful
             return results
        else:
            logger.info("CVE ID search failed or yielded no results, falling back to keyword search.")
            # Decide if we should proceed to keyword search even if a CVE ID was given but not found
            # For now, let's proceed.

    # --- Fallback to Keyword/Component Search ---
    # Ensure 'component' is treated as a list for keyword search
    if 'component' in entities and not isinstance(entities['component'], list):
        entities['component'] = [entities['component']] # Wrap single string in list
    elif 'component' not in entities:
        entities['component'] = [] # Ensure it exists as an empty list

    # Check if there are any keywords to search by
    if not entities.get('component') and not entities.get('version'):
        logger.warning("No CVE ID provided and no component/version keywords found for search.")
        return []

    return _search_nvd_by_keywords(entities, api_key, difficulty)


def summarize_attack_chain(entities: dict, cve_items: list) -> str:
    """Generates a summary using CVE details and extracted entities including network context."""
    logger.info(f"Summarizing attack chain for {len(cve_items)} CVEs.")
    cve_lines = []
    top_cve_ids = [] # Store top CVE IDs for the prompt
    for entry in cve_items[:5]: # Limit summary to first 5 CVEs
        try:
            cid = entry.id
            # Find English description
            desc_en = "No English description available."
            if hasattr(entry, 'descriptions') and entry.descriptions:
                for desc_obj in entry.descriptions:
                    if desc_obj.lang == 'en':
                        desc_en = desc_obj.value
                        break
            cve_lines.append(f"- {cid}: {desc_en}")
            if len(top_cve_ids) < 3: # Get top 3 CVE IDs for prompt context
                top_cve_ids.append(cid)
        except AttributeError as e:
            logger.warning(f"Could not process CVE entry {getattr(entry, 'id', 'UNKNOWN')}: {e}")
    cve_summary = "\n".join(cve_lines)

    # Extract context for the prompt
    components = entities.get('components', [])
    # Ensure network_context is a dict, defaulting to empty if None or missing
    network_context = entities.get('network_context') or {}
    difficulty = entities.get('difficulty', 'medium') # Get difficulty, default to medium
    context_str = f"Components: {components}\nNetwork Context: {network_context}\nRequested Difficulty: {difficulty}\nTop CVEs Found: {top_cve_ids}"
    logger.debug(f"Context for summarization LLM call: {context_str}\nFull CVE List (first 5):\n{cve_summary}")

    prompt = f"""
You are a security analyst providing a Red Team vs Blue Team assessment.

Scenario Context:
{context_str}

Relevant CVE Summaries (up to top 5):
{cve_summary}

Task: Analyze the scenario from both Red Team (attack) and Blue Team (defense/detection) perspectives, adjusting the complexity based on the 'Requested Difficulty'.

1.  **Red Team Perspective:**
     *   Describe a likely attack path, starting from an initial compromise point derived from the context (e.g., '{network_context.get('location', 'DMZ or external')}' context component like '{components[0] if components else 'external server'}').
     *   Explain how an attacker could leverage specific CVEs mentioned or component weaknesses to pivot towards other internal systems (like '{components[1:] if len(components) > 1 else 'the internal network'}'), potentially bypassing controls (like firewalls).
     *   **IMPORTANT: When selecting CVEs for the attack path, strictly prioritize those that align with the requested '{difficulty}' level.**
         - **For 'Novice'/'Easy':** Strongly prefer using CVEs known for very simple, public exploits (e.g., basic RCE, trivial info disclosure) *from the provided list*. **If no CVEs in the provided list seem suitable for a novice/easy attack, describe a generic attack using default credentials, weak passwords, or common misconfigurations INSTEAD of forcing the use of a listed CVE.**
         - **For 'Medium':** Use common pivoting techniques and relevant CVEs from the list that enable such pivots.
         - **For 'Hard'/'Master':** Feel free to chain less common or more complex CVEs from the list, or describe advanced TTPs that might leverage them.
     *   **When mentioning a specific CVE ID (e.g., CVE-XXXX-YYYYY), immediately follow it with bullet points explaining in detail:**
         - What the vulnerability fundamentally allows (e.g., RCE, info disclosure, privilege escalation, authentication bypass).
         - **Explicitly state how exploiting this vulnerability in *this step* directly enables the attacker to move to the *next system/component* mentioned in the attack path.** (e.g., "This RCE on Server A allows the attacker to plant a backdoor, which is then used to scan and connect to Server B.", "By reading credentials from file X using this info disclosure, the attacker gains valid user credentials to log into Server Y.").
     *   Mention potential tools or techniques appropriate for the '{difficulty}' level.

2.  **Blue Team Perspective:**
     *   How could this specific `{difficulty}` level attack path be *detected*? Assume tools available match the difficulty level.
         - Novice/Easy: **Assume NO SIEM or EDR is available.** Focus solely on fundamental, manual checks. Suggest specific built-in OS logs (e.g., Windows Event Logs: Security for logins ID 4625/4624, System for service errors, Application for app-specific issues; Linux: /var/log/auth.log for logins, /var/log/syslog for system messages), key application log paths. Suggest specific commands (e.g., `tasklist` or `ps` for processes, `netstat -ano` or `ss -tulnp` for connections, checking file timestamps/hashes on critical configs). **Explain *what* to look for (e.g., repeated failed logins, unexpected processes running as privileged users, connections to unknown IPs, recent config file changes).**
         - Medium: Suggest specific Event IDs (e.g., Windows Event ID 4688 for process creation), basic SIEM correlation rules (e.g., login failure spike followed by success from same IP), relevant log file patterns to `grep` for (e.g., `grep 'Accepted password for' /var/log/auth.log`). Assume basic SIEM/log aggregation might be present.
         - Hard/Master: Suggest specific EDR alert names/types (e.g., 'Credential Dumping Detected', 'Process Injection'), advanced SIEM query logic (e.g., detecting lateral movement patterns), specific network IDS signatures, or YARA rules for malware detection. Assume mature tooling.
     *   Tailor detection focus to the attack path (e.g., basic RCE attempts vs. stealthy lateral movement detection).
     *   What immediate *containment* actions should be taken if this activity is detected (e.g., isolating host '{components[0] if components else 'compromised server'}', blocking C2 IPs)?
     *   What specific *remediation* steps (beyond generic patching) are most relevant to *prevent* this scenario? Tailor the complexity and assume tools available match the difficulty level:
         - Novice/Easy: **Assume NO centralized management or advanced security tools.** Focus on essential, manual hardening steps. Suggest disabling unused services/features, enforcing strong local passwords (and changing defaults!), ensuring basic host firewall rules are configured (e.g., Windows Firewall block rules), applying critical OS/App patches manually or via basic tools. **Emphasize these are the primary defenses without advanced tooling.**
         - Medium: Suggest specific configuration settings (e.g., Group Policy settings for Windows, `sysctl` settings for Linux), enabling specific OS security features (e.g., Windows Defender Credential Guard), reviewing firewall rules between network zones. Assume some centralized configuration is possible.
         - Hard/Master: Suggest advanced hardening (e.g., application allow-listing, stricter network segmentation/micro-segmentation, implementing MFA everywhere, detailed security configuration baselines deployed via automation, deploying deception technology). Assume mature security posture management.

Output the analysis clearly structured under 'Red Team Perspective' and 'Blue Team Perspective'. Be concise and actionable.
"""
    resp = client.chat.completions.create(
        model="gpt-4.1",
        messages=[
            {"role": "system", "content": "You are a cybersecurity advisor."},
            {"role": "user", "content": prompt}
        ],
        temperature=0.7,
    )
    summary = resp.choices[0].message.content.strip()
    logger.info("Attack chain summary generated.")
    return summary

def print_welcome_message():
    """Prints a welcome banner with placeholders."""
    # ANSI escape codes for colors (e.g., Blue text)
    blue = "\033[94m"
    cyan = "\033[96m"
    reset = "\033[0m"

    # Placeholder for ASCII art - replace the string below
    ascii_art_placeholder = f"""
{blue}
           ,''',
         .' ., .',                                  ../'''',
        .'. %%, %.',                            .,/' .,%   :
       .'.% %%%,`%%%'.    .....,,,,,,.....   .,%%% .,%%'. .'
       : %%% %%%%%%',:%%>>%>' .,>>%>>%>%>>%%>,.   `%%%',% :
       : %%%%%%%'.,>>>%'   .,%>%>%'.,>%>%' . `%>>>,. `%%%:'
       ` %%%%'.,>>%'  .,%>>%>%' .,%>%>%' .>>%,. `%%>>,. `%
        `%'.,>>>%'.,%%%%%%%' .,%%>%%>%' >>%%>%>>%.`%% %% `,
        ,`%% %%>>>%%%>>%%>%%>>%>>%>%%%  %%>%%>%%>>%>%%%' % %,
      ,%>%'.>>%>%'%>>%%>%%%%>%'                 `%>%>>%%.`%>>%.
    ,%%>' .>%>%'.%>%>>%%%>>%' ,%%>>%%>%>>%>>%>%%,.`%%%>%%. `%>%.
   ` ,%' .>%%%'.%>%>>%' .,%%%%%%%%'          `%%%%.`%%%%%.%%%%> %%>%.
  ,%>%' >>%%  >%' `%%%%'     `%%%%%%%'.,>,. `%%%%'     `%%%>>%%>%
.%%>%' .%%>'  %>>%, %% oO ~ Oo %%%>>'.>>>>>>. `% oO ~ Oo'.%%%'%>%,
%>'%> .%>%>%  %%>%%%'  `OoooO'.%%>>'.>>>%>>%>>.`%`OoooO'.%%>% '%>%
%',%' %>%>%'  %>%>%>% .%,>,>,   `>'.>>%>%%>>>%>.`%,>,>' %%%%> .>%>,
` %>% `%>>%%. `%%% %' >%%%%%%>,  ' >>%>>%%%>%>>> >>%%' ,%%>%'.%%>>%.
 .%%'  %%%%>%.   `>%%. %>%%>>>%.>> >>>%>%%%%>%>>.>>>'.>%>%>' %>>%>%%
 `.%%  `%>>%%>    %%>%  %>>>%%%>>'.>%>>>>%%%>>%>>.>',%>>%'  ,>%'>% '
  %>'  %%%%%%'    `%%'  %%%%%> >' >>>>%>>%%>>%>>%> %%>%>' .%>%% .%%
 %>%>, %>%%>>%%,  %>%>% `%%  %>>  >>>%>>>%%>>>>%>>  %%>>,%>%%'.%>%,
%>%>%%, `%>%%>%>%, %>%%> ,%>%>>>.>>`.,.  `"   ..'>.%. % %>%>%'.%>%%;
%'`%%>%  %%>%%  %>% %'.>%>>%>%%>>%::.  `,   /' ,%>>>%>. >%>%'.%>%'%'
` .%>%'  >%%% %>%%'.>%>%;''.,>>%%>%%::.  ..'.,%>>%>%>,`%  %'.>%%' '
  %>%>%% `%>  >%%'.%%>%>>%>%>%>>>%>%>>%,,::,%>>%%>%>>%>%% `>>%>'
  %'`%%>%>>%  %>'.%>>%>%>>;'' ..,,%>%>%%/::%>%%>>%%,,.``% .%>%%
  `    `%>%>>%%' %>%%>>%>>%>%>%>%%>%/'       `%>%%>%>>%%% ' .%'
        %'  `%>% `%>%%;'' .,>>%>%/',;;;;;,;;;;,`%>%>%,`%'   '
        `    `  ` `%>%%%>%%>%%;/ @a;;;;;;;;;;;a@  >%>%%'
                   `/////////';, `@a@@a@@a@@aa@',;`//'
                      `//////.;;,,............,,;;//'
                          `////;;;;;;;;;;;;;;;;;/'
                             `/////////////////'      
                                      
{reset}
"""

    # Welcome Message
    welcome_text = f"{cyan}Welcome to the CVE cli! Please tell me the topology of your system and I'll help you find the vulnerabilities!{reset}"

    # Placeholder for your info - replace these
    your_name = "[j4eva]"
    contact_info = "[lfarchio@uccs.edu]"
    creator_info = f"Created by: {your_name} ({contact_info})"

    print(ascii_art_placeholder)
    print(f"{welcome_text}\n")
    print(f"{creator_info}\n")

    # Add program description
    program_info = f"""
{cyan}This tool allows you to query the NVD using natural language.
It leverages OpenAI to: Feel free to use any provider!
  - Extract key entities (components, versions, context, difficulty) from your question.
  - Search the NVD for relevant CVEs.
  - Generate Red Team vs Blue Team analysis tailored to the specified difficulty.
Type your security question or '/exit' to quit. Made with love, Luke Farch <3{reset}
"""
    print(program_info)

def main():
    nvd_key = load_api_keys()
    console = Console()
    print_welcome_message() # Print the welcome banner
    logger.info("CVE explorer. Type /exit to quit.")
    while True:
        try:
            q = console.input(">> ") # Removed undefined color variables
        except EOFError: # Handle Ctrl+D
            logger.info("EOF received. Exiting.")
            break
        except KeyboardInterrupt: # Handle Ctrl+C
             logger.info("Keyboard interrupt received. Exiting.")
             break


        if q.strip().lower() == "/exit":
            logger.info("Exit command received. Goodbye.")
            break

        console.print(Panel(Text(q, style="bold yellow"), title="Your Question", border_style="blue"))

        logger.info(f"Received query: '{q}'")
        logger.info("Parsing query...")
        entities = extract_entities(q)

        if entities:
            logger.info("Querying NVD...")
            cves = query_nvd(entities, nvd_key)
            logger.info(f"NVD query returned {len(cves)} results.")

            logger.info("Generating analysis...")
            analysis = summarize_attack_chain(entities, cves)

            console.print(Panel(
                Markdown(analysis), # Wrap analysis string with Markdown
                title="[bold green]AI Analysis[/bold green]",
                border_style="green",
                expand=False
            ))

        else:
            error_message = "Could not parse your question effectively. Please try rephrasing."
            logger.warning(error_message)
            console.print(Panel(error_message, title="Error", border_style="red"))


if __name__ == "__main__":
    main()