# from github_searcher import GithubSearcher
# from google_searcher import GoogleSearcher
import tiktoken
from llama_index.core.callbacks import CallbackManager, TokenCountingHandler
from aipt_v2.utils.doc_handler import DocHandler
from llama_index.llms.openai import OpenAI
from llama_index.llms.langchain import LangChainLLM
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
from llama_index.core import Settings
import subprocess
import shlex
import os
import openai
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import pandas as pd

# Security: CVE ID validation pattern (CWE-78 prevention)
CVE_PATTERN = re.compile(r"^CVE-\d{4}-\d{4,7}$", re.IGNORECASE)

def _validate_cve_id(cve: str) -> str:
    """
    Validate CVE ID format to prevent command injection.

    Args:
        cve: CVE identifier string

    Returns:
        Validated CVE ID

    Raises:
        ValueError: If CVE ID format is invalid
    """
    cve = cve.strip()
    if not CVE_PATTERN.match(cve):
        raise ValueError(f"Invalid CVE ID format: {cve}. Expected format: CVE-YYYY-NNNNN")
    return cve.upper()

def _sanitize_product_name(product: str) -> str:
    """
    Sanitize product name to prevent command injection.

    Args:
        product: Product name string

    Returns:
        Sanitized product name
    """
    # Remove dangerous characters that could enable shell injection
    dangerous_chars = [";", "&", "|", "$", "`", "\n", "\r", "\\", "'", '"', "(", ")", "{", "}", "[", "]", "<", ">"]
    sanitized = product
    for char in dangerous_chars:
        sanitized = sanitized.replace(char, "")
    return sanitized.strip()[:200]  # Limit length

from sklearn.metrics import confusion_matrix, cohen_kappa_score
from sklearn.preprocessing import MinMaxScaler, StandardScaler, RobustScaler
import seaborn as sns
import matplotlib.pyplot as plt
from scipy.stats import pearsonr, spearmanr
from itertools import combinations
import numpy as np
import json
import logging
import time
from tqdm import tqdm
from aipt_v2.utils.searchers.search_once import compose
import re
import yaml
from aipt_v2.utils.model_manager import get_model
logger = logging.getLogger(__name__)

config_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'configs', 'config.yaml')
with open(config_path, 'r', encoding='utf-8') as f:
    config = yaml.safe_load(f)
planning_config = config['runtime']['planning']
model_name_for_token = config['models']['openai']['model']

def cvemap_search(cve, info_dir):
    """
    Search CVE information using cvemap tool.

    Security: Uses validated CVE ID and subprocess without shell=True (CWE-78 fix).
    """
    # Validate CVE ID to prevent command injection
    validated_cve = _validate_cve_id(cve)
    cvemap_json_path = f"{info_dir}/cvemap.json"

    try:
        # SECURE: Use subprocess without shell=True
        # Run cvemap with -id flag instead of piping through echo
        result = subprocess.run(
            ["cvemap", "-id", validated_cve, "-json"],
            capture_output=True,
            text=True,
            timeout=50,
            check=False  # Handle errors ourselves
        )

        if result.returncode != 0:
            logging.error(f"cvemap failed for {validated_cve}: {result.stderr}")
            return None

        # Write output to file
        with open(cvemap_json_path, 'w') as f:
            f.write(result.stdout)

        json_data = json.loads(result.stdout)
        if len(json_data) == 0:
            logging.error(f"Error in cvemap search: {validated_cve}")
            return None
        return json_data[0]

    except subprocess.TimeoutExpired:
        logging.error(f"cvemap search timed out for {validated_cve}")
        return None
    except json.JSONDecodeError as e:
        logging.error(f"Invalid JSON from cvemap for {validated_cve}: {e}")
        return None
    except FileNotFoundError:
        logging.error("cvemap command not found. Please install cvemap.")
        return None

def searchsploit_search(cve):
    """
    Search ExploitDB using searchsploit tool.

    Security: Uses validated CVE ID and subprocess without shell=True (CWE-78 fix).
    """
    # Validate CVE ID to prevent command injection
    validated_cve = _validate_cve_id(cve)
    searchsploit_json_path = f"resources/{validated_cve}/exploitdb.json"

    # Ensure directory exists
    os.makedirs(os.path.dirname(searchsploit_json_path), exist_ok=True)

    try:
        # SECURE: Use subprocess without shell=True
        result = subprocess.run(
            ["searchsploit", validated_cve, "-j"],
            capture_output=True,
            text=True,
            timeout=50,
            check=False
        )

        # Write output to file
        with open(searchsploit_json_path, 'w') as f:
            f.write(result.stdout)

        json_data = json.loads(result.stdout)
        return json_data

    except subprocess.TimeoutExpired:
        logging.error(f"searchsploit search timed out for {validated_cve}")
        return None
    except json.JSONDecodeError as e:
        logging.error(f"Invalid JSON from searchsploit for {validated_cve}: {e}")
        return None
    except FileNotFoundError:
        logging.error("searchsploit command not found. Please install exploitdb.")
        return None
    

def categorize_cvss(cvss_score):
    if not (0.0 <= cvss_score <= 10.0):
        return "value not in range"

    if cvss_score < 4.0:
        return "hard"
    elif 4.0 <= cvss_score < 7.0:
        return "medium"
    else:
        return "easy"
    
def categorize_epss(epss_score):
    if not (0.0 <= epss_score <= 1.0):
        return "value not in range"

    if epss_score < 0.4:
        return "hard"
    elif 0.4 <= epss_score < 0.94:
        return "medium"
    else:
        return "easy"
    
def count_cwe(cve_lst):
    
    cwe_dict = {}
    for cve in cve_lst:
        info_dir = f"resources/{cve}/info"
        cvemap_json_path = f"{info_dir}/cvemap.json"
        with open(cvemap_json_path) as f:
            cvemap_json = json.load(f)[0]
            cwe_lst = cvemap_json.get('weaknesses', [])
            for cwe in cwe_lst:
                cwe_id = cwe['cwe_id']
                if cwe_id in cwe_dict:
                    cwe_dict[cwe_id] += 1
                else:
                    cwe_dict[cwe_id] = 1
    return cwe_dict

def calculate_score(features, trending_score):
    weights = config['cve_scoring']['weights']

    scores = {}
    max_score_github = 0
    max_score_repo_github = None
    max_score_expdb = 0
    max_score_repo_expdb = None
    max_score = 0
    max_score_repo = None
    final_score = 0
    has_code = False

    code_sources = ["GitHub", "ExploitDB"]

    code_results = None
    doc_results = None

    if features['code'].get('GitHub') or features['code'].get('ExploitDB'):
        for source in code_sources:
            if not features['code'].get(source):
                continue  # if not have data from source, skip

            has_code = True
            for repo, vul_type in features['code'][source]["vul_type"].items():
                score = 0
                
                # Calculate scores for each field
                score += weights["vul_type"].get(vul_type, 0)
                score += weights["isRemote"].get(features['code'][source]["isRemote"].get(repo, ""), 0)
                
                # Attack complexity fields
                attack_complexity = features['code'][source]["attack_complexity"].get(repo, {})
                for field, value in attack_complexity.items():
                    score += 1 * weights["attack_complexity"].get(field, {}).get(value, 0)
                
                if repo == "Code_File":
                    score = score / 2

                score *= weights["exp_maturity"].get(features['code'][source]["exp_maturity"].get(repo, ""), 0)

                score *= weights["lang_class"].get(features['code'][source]["lang_class"].get(repo, ""), 1)

                if source == "GitHub":
                    score *= weights["source_weights"]["gthb"]
                elif source == "ExploitDB":
                    score *= weights["source_weights"]["expdb"]

                # Store the score
                scores[repo] = score

                # Check for max score
                if repo.isdigit():  # ExploitDB
                    if score >= max_score_expdb:
                        max_score_expdb = score
                        max_score_repo_expdb = repo
                else:  # GitHub
                    if score >= max_score_github:
                        max_score_github = score
                        max_score_repo_github = repo

        if max_score_expdb > max_score_github:
            max_score_repo = max_score_repo_expdb
            max_score = max_score_expdb
        else:
            max_score_repo = max_score_repo_github
            max_score = max_score_github

        if has_code:
            sorted_scores = sorted(scores.items(), key=lambda item: item[1], reverse=True)

            if trending_score == 999:
                trending_score_github = 0
                trending_score_expdb = weights["expdb_default_score"]
            else:
                trending_score_github = trending_score
                if features['code'].get('ExploitDB'):
                    trending_score_expdb = weights["expdb_default_score"]
                else: trending_score_expdb = 0

            # Add weighted trending score to the max score
            trend_score_weighted_expdb = trending_score_expdb * weights["trending_score"]
            trend_score_weighted_github = trending_score_github * weights["trending_score"]
    
            # Handle cases where one source has no scores
            final_expdb = (max_score_expdb + trend_score_weighted_expdb) if max_score_expdb > 0 else 0
            final_github = (max_score_github + trend_score_weighted_github) if max_score_github > 0 else 0

            final_score = max(final_expdb, final_github)

            code_results = (sorted_scores, max_score_repo, max_score, final_score, has_code)
    
    if features['doc']:
        doc_score = 0
        doc_score += weights["vul_type"].get(features["doc"]["vul_type"], 0)
        doc_score += weights["isRemote"].get(features["doc"]["isRemote"], 0)

        # Attack complexity fields for 'doc'
        for field, value in features["doc"]["attack_complexity"].items():
            doc_score += weights["attack_complexity"].get(field, {}).get(value, 0)

        doc_score *= weights["source_weights"]["gg"]
        
        scores["doc"] = doc_score

        # Check if 'doc' has the max score
        if doc_score > max_score:
            max_score = doc_score
            max_score_repo = "doc"
            final_score = doc_score

        doc_results = (scores, max_score_repo, max_score, max_score, has_code)
    
    if code_results and doc_results:
        code_sorted_scores, code_max_repo, code_max_score, code_final_score, code_has_code = code_results
        doc_scores, doc_max_repo, doc_max_score, doc_final_score, doc_has_code = doc_results
        
        # merge scores
        all_scores = dict(code_sorted_scores)
        all_scores.update(doc_scores)
        sorted_all_scores = sorted(all_scores.items(), key=lambda item: item[1], reverse=True)
        
        # final max score and response repo
        if doc_max_score > code_max_score:
            final_max_score = doc_max_score
            final_max_repo = doc_max_repo
            final_score = doc_max_score
        else:
            final_max_score = code_max_score
            final_max_repo = code_max_repo
            final_score = code_final_score
            
        return sorted_all_scores, final_max_repo, final_max_score, final_score, has_code or True
        
    elif code_results:
        return code_results
    elif doc_results:
        return doc_results
    else:
        return None, None, max_score, final_score, has_code
    
def calculate_match_stats(values1, values2):
    total = len(values1)
    exact_matches = sum(v1 == v2 for v1, v2 in zip(values1, values2))
    near_matches = sum(
        abs(['easy', 'medium', 'hard'].index(v1) - ['easy', 'medium', 'hard'].index(v2)) == 1
        for v1, v2 in zip(values1, values2)
    )
    mismatches = total - exact_matches - near_matches
    return {
        "Exact Match (%)": (exact_matches / total) * 100,
        "Near Match (%)": (near_matches / total) * 100,
        "Mismatch (%)": (mismatches / total) * 100,
    }

def create_df(cvss_data, epss_data, pentestasst_data):
    all_keys = sorted(set(cvss_data.keys()).union(epss_data.keys(), pentestasst_data.keys()))
    data = {
        'CVSS': [cvss_data.get(key, np.nan) for key in all_keys],
        'EPSS': [epss_data.get(key, np.nan) for key in all_keys],
        'EEAS': [pentestasst_data.get(key, np.nan) for key in all_keys],
    }
    df = pd.DataFrame(data, index=all_keys)
    return df

def normalize_data(df):
    if len(df) == 1:
        # for the case of a single value, directly return 0.5 or keep the original value
        df[['EEAS']] = 0.5  # or df[['EEAS']]
    else:
        scaler = StandardScaler()
        minmax_scaler = MinMaxScaler()
        df[['EEAS']] = minmax_scaler.fit_transform(scaler.fit_transform(df[['EEAS']]))
    normalized_df = pd.DataFrame(df, index=df.index, columns=df.columns)
    return normalized_df
    
def bin_agreement_analysis(df):
    pairwise_stats = {}
    for (name1, col1), (name2, col2) in combinations(df.items(), 2):
        # Calculate confusion matrix
        cm = confusion_matrix(col1, col2, labels=['easy', 'medium', 'hard'])
        kappa = cohen_kappa_score(col1, col2, labels=['easy', 'medium', 'hard'])
        stats = calculate_match_stats(col1, col2)

        pairwise_stats[f"{name1} vs {name2}"] = {
            "Confusion Matrix": cm.tolist(),
            "CohenS Kappa": kappa,
            **stats,
        }

    # Display results
    for pair, results in pairwise_stats.items():
        print(f"\nPairwise comparison: {pair}")
        print(f"Confusion Matrix:\n{np.array(results['Confusion Matrix'])}")
        print(f"Cohen's Kappa: {results['CohenS Kappa']:.4f}")
        print(f"Exact Match (%): {results['Exact Match (%)']:.2f}")
        print(f"Near Match (%): {results['Near Match (%)']:.2f}")
        print(f"Mismatch (%): {results['Mismatch (%)']:.2f}")

    return df

def num_agreement_analysis(df):
    stats = {}
    for col1, col2 in combinations(df.columns, 2):
        mean_values = (df[col1] + df[col2]) / 2
        diff_values = df[col1] - df[col2]

        mean_diff = np.mean(diff_values)
        std_diff = np.std(diff_values)

        print(f"diff mean: {np.mean(diff_values)}")
        print(f"diff std: {np.std(diff_values)}")

        threshold = 1.96 * std_diff
        diff_condition = abs(df[col1] - df[col2]) > threshold

        rows_large_diff = df[diff_condition]
        print(f"Rows where the difference between {col1} and {col2} is larger than {threshold}:")
        print(rows_large_diff)
        print(f"Total rows: {len(df)}, Rows with large difference: {len(rows_large_diff)}, percentage: {len(rows_large_diff)/len(df) * 100}%")
        
        values1, values2 = df[col1], df[col2]
        pearson_corr, _ = pearsonr(values1, values2)
        spearman_corr, _ = spearmanr(values1, values2)
        mae = np.mean(np.abs(values1 - values2))
        rmse = np.sqrt(np.mean((values1 - values2) ** 2))
        stats[f"{col1} vs {col2}"] = {
            "Pearson Correlation": pearson_corr,
            "Spearman Correlation": spearman_corr,
            "Mean Absolute Error (MAE)": mae,
            "Root Mean Squared Error (RMSE)": rmse,
        }

    for pair, metrics in stats.items():
        print(f"\nAgreement Analysis for {pair}:")
        for metric, value in metrics.items():
            print(f"{metric}: {value:.4f}")

    return df

def visualize_num_results(df, output_dir="plots/"):

    plt.rc('font', size=14) 

    # Pairwise scatter plots
    pairplot = sns.pairplot(df, kind="reg", diag_kind="kde")
    pairplot.savefig(os.path.join(output_dir, "pairwise_scatter_plots.pdf"))
    plt.close()

    # Heatmap for correlation
    correlation_matrix = df.corr()
    plt.figure(figsize=(8, 6))
    sns.heatmap(correlation_matrix, annot=True, cmap="coolwarm", fmt=".2f")
    plt.title("Heatmap of Correlations")
    plt.savefig(os.path.join(output_dir, "correlation_heatmap.pdf"))
    plt.close()

    # Bland-Altman plots
    for col1, col2 in combinations(df.columns, 2):
        mean_values = (df[col1] + df[col2]) / 2
        diff_values = df[col1] - df[col2]

        plt.figure(figsize=(8, 6))
        plt.scatter(mean_values, diff_values, alpha=0.7)
        plt.axhline(np.mean(diff_values), color='red', linestyle='--', label='Mean Difference')
        plt.axhline(np.mean(diff_values) + 1.96 * np.std(diff_values), color='blue', linestyle='--', label='±1.96SD')
        plt.axhline(np.mean(diff_values) - 1.96 * np.std(diff_values), color='blue', linestyle='--', label=None)
        plt.title(f"Bland-Altman Plot: {col1} vs {col2}", weight='bold')
        plt.xlabel('Mean of Two Measurements', weight='bold')
        plt.ylabel('Difference', weight='bold')
        plt.legend()
        filename = f"bland_altman_{col1}_vs_{col2}.pdf"
        plt.savefig(os.path.join(output_dir, filename))
        plt.close()

def visualize_bin_results(cvss_data, epss_data, pentestasst_data, output_dir):
    # Confusion Matrix Heatmap
    for (name1, data1), (name2, data2) in combinations([("CVSS", cvss_data), ("EPSS", epss_data), ("EEAS", pentestasst_data)], 2):
        labels = ['easy', 'medium', 'hard']
        confusion = confusion_matrix(list(data1.values()), list(data2.values()), labels=labels)
        
        plt.figure(figsize=(6, 5))
        sns.heatmap(confusion, annot=True, fmt="d", cmap="Blues", xticklabels=labels, yticklabels=labels)
        plt.title(f"Confusion Matrix: {name1} vs {name2}")
        plt.xlabel(f"{name2} Predictions")
        plt.ylabel(f"{name1} Predictions")
        plt.savefig(f"{output_dir}/confusion_matrix_{name1}_vs_{name2}.pdf")
        plt.close()

    # Bar Chart of Class Distributions
    data = {
        "Category": ['easy', 'medium', 'hard'],
        "CVSS": [list(cvss_data.values()).count(c) for c in ['easy', 'medium', 'hard']],
        "EPSS": [list(epss_data.values()).count(c) for c in ['easy', 'medium', 'hard']],
        "EEAS": [list(pentestasst_data.values()).count(c) for c in ['easy', 'medium', 'hard']],
    }
    df = pd.DataFrame(data).set_index("Category")
    
    df.plot(kind="bar", figsize=(8, 6))
    plt.title("Class Distributions Across Dictionaries")
    plt.ylabel("Frequency")
    plt.xlabel("Category")
    plt.xticks(rotation=0)
    plt.savefig(f"{output_dir}/class_distributions.pdf")
    plt.close()

    # Pairwise Agreement Bar Chart
    pairwise_agreements = []
    for (name1, data1), (name2, data2) in combinations([("CVSS", cvss_data), ("EPSS", epss_data), ("EEAS", pentestasst_data)], 2):
        total = len(data1)
        exact_matches = sum(1 for k in data1 if data1[k] == data2[k])
        near_matches = sum(1 for k in data1 if abs(['easy', 'medium', 'hard'].index(data1[k]) -
                                                   ['easy', 'medium', 'hard'].index(data2[k])) == 1)
        mismatches = total - exact_matches - near_matches

        pairwise_agreements.append({
            "Pair": f"{name1} vs {name2}",
            "Exact Matches (%)": exact_matches / total * 100,
            "Near Matches (%)": near_matches / total * 100,
            "Mismatches (%)": mismatches / total * 100
        })

    pairwise_df = pd.DataFrame(pairwise_agreements).set_index("Pair")
    pairwise_df.plot(kind="bar", stacked=True, figsize=(10, 6))
    plt.title("Pairwise Agreement Percentages")
    plt.ylabel("Percentage")
    plt.xlabel("Pairwise Comparison")
    plt.xticks(rotation=0)
    plt.legend(loc="upper right")
    plt.savefig(f"{output_dir}/pairwise_agreements.pdf")
    plt.close()
    
def cve_classifier(cve, output_dir="resources/", mode = "specific"):
    trending_score_path = f"{output_dir}/{cve}/Trend_Score.json"
    feature_path = f"{output_dir}/{cve}/features.json"
    cvemap_path = f"{output_dir}/{cve}/info/cvemap.json"
    result = None
    with open(feature_path, 'r') as f:
        result = json.load(f)
    trending_score = 0
    if result['code'].get('GitHub'):
        try:
            with open(trending_score_path, 'r') as f:
                _trending_score = json.load(f)
                # print(trending_score)
                trending_score = min(_trending_score.get("trend_score", 0), 50)
        except:
            logging.warning(f"Trending score not found for {cve}")
    else:
        trending_score = 999

    if mode == "specific":
        cvemap_json = None
        with open(cvemap_path, 'r') as f:
            cvemap_json = json.load(f)[0]
        cvss_score = cvemap_json.get('cvss_score', 0)
        epss_score = cvemap_json['epss']['epss_score']
        epss_percentile = cvemap_json['epss']['epss_percentile']
    elif mode == "general":
        cvss_score = 0
        epss_score = 0
        epss_percentile = 0
    cvss_category = categorize_cvss(cvss_score)
    epss_category = categorize_epss(epss_percentile)

    scores, max_score_repo, max_score, final_score, has_code = calculate_score(result, trending_score)

    exploitability = "hard"
    if final_score is not None and final_score > 0:
        if final_score > 50:    
            exploitability = "easy"
        elif final_score > 35:
            exploitability = "medium"
        else:
            exploitability = "hard"
    

    with open(f"{output_dir}/{cve}/classification.json", 'w') as f:
        classification = {
            "cvss_category": cvss_category,
            "epss_category": epss_category,
            "exploitability": exploitability,
            "final_score": final_score,
            "max_score_repo": max_score_repo,
            "max_score": max_score,
            "scores": scores,
            "has_code": has_code
        }
        json.dump(classification, f, indent=4)

    return cvss_score, cvss_category, epss_percentile, epss_category, final_score, exploitability, has_code

def cve_analysis(cve, output_dir="resources/"):
    cve_dir = f"{output_dir}/{cve}"
    info_dir = f"{output_dir}/{cve}/info"
    searching_start_time = time.time()
    compose(output_dir, cve)
    if not os.path.exists(info_dir):
        os.mkdir(info_dir)
    cvemap_json = cvemap_search(cve, info_dir)
    if cvemap_json is None:
        return 0
    cve_description = cvemap_json['cve_description']

    logging.info(f"Analyzing {cve}")
    doc_handler = DocHandler()

    analysis_start_time = time.time()
    result = doc_handler.vul_analysis(cve, output_dir, cve_description)
    analysis_end_time = time.time()
    searching_time = analysis_start_time - searching_start_time
    analysis_time = analysis_end_time - analysis_start_time
    logging.info(f"Analysis time: {analysis_time} seconds")
    
    with open(f"{output_dir}/{cve}/features.json", 'w') as f:
        
        json.dump(result, f, indent=4)

    return searching_time, analysis_time

def general_analysis(keyword, output_dir="resources/"):
    searching_start_time = time.time()
    compose(output_dir, keyword, loose_mode = True)

    cve = keyword

    cve_description = ""

    logging.info(f"Analyzing {keyword}")
    doc_handler = DocHandler()

    analysis_start_time = time.time()
    result = doc_handler.vul_analysis(cve, output_dir, cve_description)
    analysis_end_time = time.time()
    searching_time = analysis_start_time - searching_start_time
    analysis_time = analysis_end_time - analysis_start_time
    logging.info(f"Analysis time: {analysis_time} seconds")
    
    with open(f"{output_dir}/features.json", 'w') as f:
        
        json.dump(result, f, indent=4)

    return searching_time, analysis_time

def cve_analysis_from_epss_csv(csv_path):
    
    df = pd.read_csv(csv_path)
    # print(df.cve)
    with tqdm(total=len(df.index), desc=f'Analyzing CVEs') as pbar:
        for index, row in df.iterrows():
            cve = row['cve']
            if int(cve.split('-')[1]) < 2017 or int(cve.split('-')[1]) > 2022:
                logger.info(f"Skipping {cve}")
                pbar.update()
                continue
            try:
                cve_analysis(cve)
                time.sleep(1)
                pbar.update()
                logging.info(f"Finished {index}: {cve}")
            except Exception as e:
                logging.error(f"Error in {cve}: {e}")
                pbar.update()
                continue

def product_to_cve(product, output_dir):
    """
    Search for CVEs related to a product using cvemap.

    Security: Uses sanitized product name and subprocess without shell=True (CWE-78 fix).
    """
    cve_lst = []
    product_dir_name = product.lower().replace(" ", "_").replace("/", "_").replace(":", "_").replace("\\", "_").replace("(", "_").replace(")", "_").replace('"', "").replace("'", "").replace("\n", "").replace("&", "")
    if os.path.exists(f"{output_dir}/{product_dir_name}"):
        with open(f"{output_dir}/{product_dir_name}/cve_lst.json") as f:
            cve_lst = json.load(f)
        return cve_lst

    os.makedirs(f"{output_dir}/{product_dir_name}", exist_ok=True)
    cleaned_product = product_keyword_gen_openai(product, 3)
    # print(cleaned_product)
    cleaned_product.append(product_dir_name)

    for p in cleaned_product:
        # Security: Sanitize product name to prevent command injection
        sanitized_product = _sanitize_product_name(p)
        if not sanitized_product:
            logging.warning(f"Skipping empty product name after sanitization: {p}")
            continue

        p_dir_name = p.lower().replace(" ", "_").replace("/", "_").replace(":", "_").replace("\\", "_").replace("(", "_").replace(")", "_").replace('"', "").replace("'", "").replace("\n", "").replace("&", "")
        cvemap_json_path = f"{output_dir}/{product_dir_name}/cvemap_{p_dir_name}.json"

        try:
            # SECURE: Use subprocess without shell=True
            # Run cvemap with -p flag and capture output directly
            result = subprocess.run(
                ["cvemap", "-p", sanitized_product, "-json"],
                capture_output=True,
                text=True,
                timeout=50,
                check=False  # Handle errors ourselves
            )

            if result.returncode != 0:
                logging.warning(f"cvemap failed for product {sanitized_product}: {result.stderr}")
                continue

            # Write output to file (instead of shell redirection)
            with open(cvemap_json_path, 'w') as f:
                f.write(result.stdout)

            json_data = json.loads(result.stdout)
            for cve in json_data:
                cve_lst.append(cve['cve_id'])

        except subprocess.TimeoutExpired:
            logging.error(f"cvemap search timed out for product {sanitized_product}")
            continue
        except json.JSONDecodeError as e:
            logging.warning(f"Invalid JSON from cvemap for product {sanitized_product}: {e}")
            continue
        except FileNotFoundError:
            logging.error("cvemap command not found. Please install cvemap.")
            break

        with open(f"{output_dir}/{product_dir_name}/cve_lst.json", 'w') as f:
            json.dump(cve_lst, f, indent=4)

    return cve_lst

def product_keyword_gen_openai(product, num):
    prompt = (
        "You're an excellent system administrator. You will be given a product name, which may not be in a standard format. "
        "Your task is to generate alternative product names to broaden the search scope.\n"
        f"The given product name is {product}, please generate at most {num} alternative product names in the following format.\n"
        "[Alternative product name 1, Alternative product name 2, Alternative product name 3, ...]"
    )
    try:
        model_name = planning_config['model']
        llm = get_model(model_name)
        from langchain_core.messages import HumanMessage
        response = llm.invoke([HumanMessage(content=prompt)])
        product_lst_str = response.content.strip() if hasattr(response, 'content') else str(response)
    except Exception as e:
        import openai
        client = openai.OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[
                {"role": "system", "content": prompt},
            ],
        )
        product_lst_str = str(response.choices[0].message.content)
    product_lst = product_lst_str.strip("[]").replace("'", "").split(", ")
    # print(product_lst)
    product_lst = [product.lower().replace(" ", "_").replace("/", "_").replace(":", "_").replace("\\", "_").replace("(", "_").replace(")", "_").replace('"', "").replace("'", "").replace("\n", "").replace("&", "") for product in product_lst]

    return product_lst

def product_keyword_gen_huggingface(product, num):

    messages = [
        {"role": "system", "content": "You're an excellent system administrator. You will be given a product name, which may not be in a standard format. Your task is to generate alternative product names to broaden the search scope."},
        {"role": "user", "content": f"The given product name is {product}, please generate {num} alternative product names in the following format: [Alternative product name 1, Alternative product name 2, Alternative product name 3, ...]"},
    ]

    # Run with Local LLM
    product_lst = chat_completion_huggingface(
        model_name="meta-llama/Meta-Llama-3.1-8B-Instruct",
        messages=messages,
        max_tokens=200,
        temperature=0.0
    )

    return product_lst

def chat_completion_huggingface(
    model_name="meta-llama/Meta-Llama-3.1-8B-Instruct",
    messages=None,
    max_tokens=100,
    temperature=0.0,
    top_p=0.0,
    revision="5206a32e0bd3067aef1ce90f5528ade7d866253f"  # Security: Pinned commit hash (CWE-494)
):
    """
    Generates a chat completion using a specified Hugging Face model.

    Args:
        model_name (str): The Hugging Face model name (default is Mistral-7B-Instruct).
        messages (list): A list of chat messages in OpenAI format [{"role": "user", "content": "Hello"}].
        max_tokens (int): Maximum number of tokens to generate.
        temperature (float): Sampling temperature (higher = more creative, lower = more deterministic).
        top_p (float): Nucleus sampling parameter (controls randomness).
        revision (str): Model revision to use (commit hash or branch name for reproducibility).

    Returns:
        str: The model's response.

    Security:
        Using 'revision' parameter pins the model to a specific version,
        preventing supply chain attacks where a malicious model update
        could be downloaded automatically.
    """

    # Load the model and tokenizer dynamically
    # Security: Pin model revision to prevent supply chain attacks (CWE-494)
    # Bandit requires string literals for revision detection
    pinned_revision = revision if revision != "main" else "5206a32e0bd3067aef1ce90f5528ade7d866253f"
    tokenizer = AutoTokenizer.from_pretrained(
        model_name,
        revision=pinned_revision,  # nosec B615 - revision is pinned via function default
        trust_remote_code=False  # Security: Never execute remote code
    )
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_map="auto",  # Uses GPU if available
        revision=pinned_revision,  # nosec B615 - revision is pinned via function default
        trust_remote_code=False  # Security: Never execute remote code
    )

    # Convert messages into a formatted prompt
    prompt = ""
    for message in messages:
        if message["role"] == "system":
            prompt += f"<s>[SYSTEM]: {message['content']}\n"
        elif message["role"] == "user":
            prompt += f"[USER]: {message['content']}\n"
        elif message["role"] == "assistant":
            prompt += f"[ASSISTANT]: {message['content']}\n"
    prompt += "[ASSISTANT]:"

    # Tokenize input
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # Generate response
    output = model.generate(
        **inputs,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        do_sample=False
    )

    # Decode and return response
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response.split("[ASSISTANT]:")[-1].strip()

def analyze_cve_lst(cve_lst, output_dir, app_name):
    cvss_results = {}
    cvss_scores = {}
    epss_results = {}
    epss_scores = {}
    pentestasst_results = {}
    pentestasst_scores = {}
    analysis_time_dict = {}
    has_code_lst = []
    total_searching_time = 0
    total_analysis_time = 0

    for cve in cve_lst:
        logging.info(f"Analyzing {cve}")
        if not os.path.exists(f"{output_dir}/{cve}/features.json"):
            searching_time, analysis_time = cve_analysis(cve, output_dir)
            if analysis_time == 0:
                continue
            total_searching_time += searching_time
            total_analysis_time += analysis_time
            analysis_time_dict[cve] = analysis_time
        cvss_score, cvss_category, epss_score, epss_category, final_score, exploitability, has_code = cve_classifier(cve, output_dir, mode = "specific")
        cvss_results[cve] = cvss_category
        cvss_scores[cve] = cvss_score
        epss_results[cve] = epss_category
        epss_scores[cve] = epss_score
        pentestasst_results[cve] = exploitability
        pentestasst_scores[cve] = final_score
        if has_code:
            has_code_lst.append(cve)

    with open(f"{output_dir}/analysis_time.json", 'w') as f:
        json.dump(analysis_time_dict, f, indent=4)

    # bin_df = create_df(cvss_results, epss_results, pentestasst_results)
    num_df = create_df(cvss_scores, epss_scores, pentestasst_scores)
    normalized_num_df = normalize_data(num_df)
    normalized_has_code_num_df = normalized_num_df.loc[has_code_lst]
    num_df.to_csv(f'{output_dir}/score_results.csv')
    normalized_num_df.to_csv(f'{output_dir}/normalized_score_results.csv')
    normalized_has_code_num_df.to_csv(f'{output_dir}/normalized_has_code_score_results.csv')


    if app_name:
        logging.info(f"Searching for general exp info for {app_name}")
        print("Start searching general exp info...")
        app_exp_path = app_name + "_exp"
        if not os.path.exists(f"{output_dir}/{app_exp_path}/features.json"):
            keyword = app_name + " exploit"
            searching_time, analysis_time = general_analysis(keyword, os.path.join(output_dir, app_exp_path))
            total_searching_time += searching_time
            total_analysis_time += analysis_time
        cve_classifier(app_exp_path, output_dir, mode = "general")

    avg_searching_time = total_searching_time / ((len(cve_lst) +1) if app_name else (len(cve_lst)))
    avg_analysis_time = total_analysis_time / ((len(cve_lst) +1) if app_name else (len(cve_lst)))
    logging.info(f"Total searching time: {total_searching_time} seconds, average time: {avg_searching_time} seconds")
    logging.info(f"Total analysis time: {total_analysis_time} seconds, average time: {avg_analysis_time} seconds")

    return total_searching_time, total_analysis_time


def get_exp_info(cve_lst = [], output_dir = "", app_name = ""):
    logging.basicConfig(format='%(asctime)s %(levelname)-8s %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S',
                        filename='cve_info.log',
                        level=logging.INFO)
    
    track_tokens = True
    token_counter = None
    model_name = planning_config['model']
    if track_tokens:
        try:
            token_counter = TokenCountingHandler(tokenizer=tiktoken.encoding_for_model(model_name_for_token).encode)
            Settings.callback_manager = CallbackManager([token_counter])
        except Exception as e:
            logger.warning(f"Tokenizer for model {model_name_for_token} not found, token counting disabled: {e}")
            Settings.callback_manager = CallbackManager([])
    
    ### Step 1: Select the LLM

    ## Initialize LLM from config via model manager (provider-agnostic)
    try:
        model_name = planning_config['model']
    except Exception:
        model_name = "openai"

    llm = get_model(model_name)
    if llm is None:
        # Keep previous behavior when no API key/config: abort gracefully
        print("LLM_API_KEY not set")
        return

    try:
        Settings.llm = LangChainLLM(llm=llm)
    except Exception:
        # Fallback to OpenAI config if available to preserve behavior
        cve_config = config['cve']
        print(f"Model: {cve_config['model']}")
        Settings.llm = OpenAI(temperature=cve_config['temperature'], model=cve_config['model'])

    ### Step 2: Input the CVEs to analyze

    if not cve_lst:
        print("CVE to be searched not set! ")
        return
    

    ### Step 3: Analyze the CVEs
    if not output_dir:
        print("Directory to store not set! will use default settings.")
        output_dir = "resources_cve_lst_try"

    os.makedirs(output_dir, exist_ok=True)

    ## Scenario 2: Given a list of CVEs, analyze the CVEs

    total_searching_time, total_analysis_time = analyze_cve_lst(cve_lst, output_dir, app_name)


    if track_tokens:
        print(
            "Embedding Tokens: ",
            token_counter.total_embedding_token_count,
            "\n",
            "LLM Prompt Tokens: ",
            token_counter.prompt_llm_token_count,
            "\n",
            "LLM Completion Tokens: ",
            token_counter.completion_llm_token_count,
            "\n",
            "Total LLM Token Count: ",
            token_counter.total_llm_token_count,
            "\n",
            "Average LLM Token Count: ",
            token_counter.total_llm_token_count / len(cve_lst),
            "\n",
        )

        logging.info(
            "Embedding Tokens: "
            + str(token_counter.total_embedding_token_count)
            + "\n"
            + "LLM Prompt Tokens: "
            + str(token_counter.prompt_llm_token_count)
            + "\n"
            + "LLM Completion Tokens: "
            + str(token_counter.completion_llm_token_count)
            + "\n"
            + "Total LLM Token Count: "
            + str(token_counter.total_llm_token_count)
            + "\n"
            + "Average LLM Token Count: "
            + str(token_counter.total_llm_token_count / len(cve_lst))
            + "\n"
        )

        token_counter.reset_counts()

    return total_searching_time, total_analysis_time


def main():
    logging.basicConfig(format='%(asctime)s %(levelname)-8s %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S',
                        filename='cve_info.log',
                        level=logging.INFO)
    
    track_tokens = False
    token_counter = None
    if track_tokens:
        token_counter = TokenCountingHandler(tokenizer=tiktoken.encoding_for_model("gpt-4o-mini").encode)
        Settings.callback_manager = CallbackManager([token_counter])
    
    ### Step 1: Select the LLM

    ## OpenAI API 
    api_key = os.environ.get("OPENAI_API_KEY")
    if api_key is None or api_key == "":
        print("OPENAI_API_KEY not set")
        return
    elif "f4Ph3uIoqGLC9" in api_key:
        print("Model: o3-mini")
        Settings.llm = OpenAI(temperature=0, model="o3-mini")
    elif "klO8n1OFxLWBoPIeDycM" in api_key:
        print("Model: gpt-4o-mini")
        Settings.llm = OpenAI(temperature=0, model="gpt-4o-mini")
    else:
        print("Model: gpt-4o-mini (default)")
        Settings.llm = OpenAI(temperature=0, model="gpt-4o-mini")
    # Settings.llm = OpenAI(temperature=0, model="o3-mini")

    ## Option 2: Read CVE list from file
    cve_lst = []
    with open("../../target_lst/vulhub_test/xstream.txt", 'r') as f:
        cve_lst = f.read().splitlines()

    ## Option 3: Hard-coded CVE list
    # cve_lst = ['CVE-XXX-XXXX', 'CVE-XXX-XXXX', 'CVE-XXX-XXXX', 'CVE-XXX-XXXX', 'CVE-XXX-XXXX']
    

    ### Step 3: Analyze the CVEs
    output_dir = "resources_test/xstream"
    os.makedirs(output_dir, exist_ok=True)

    ## Scenario 1: Given a list of CVEs, analyze the CVEs

    # analyze_cve_lst(cve_lst, output_dir)
    analyze_cve_lst(cve_lst, output_dir, "xstream")  # assume app's name is xstream


    if track_tokens:
        print(
            "Embedding Tokens: ",
            token_counter.total_embedding_token_count,
            "\n",
            "LLM Prompt Tokens: ",
            token_counter.prompt_llm_token_count,
            "\n",
            "LLM Completion Tokens: ",
            token_counter.completion_llm_token_count,
            "\n",
            "Total LLM Token Count: ",
            token_counter.total_llm_token_count,
            "\n",
            "Average LLM Token Count: ",
            token_counter.total_llm_token_count / len(cve_lst),
            "\n",
        )

        logging.info(
            "Embedding Tokens: "
            + str(token_counter.total_embedding_token_count)
            + "\n"
            + "LLM Prompt Tokens: "
            + str(token_counter.prompt_llm_token_count)
            + "\n"
            + "LLM Completion Tokens: "
            + str(token_counter.completion_llm_token_count)
            + "\n"
            + "Total LLM Token Count: "
            + str(token_counter.total_llm_token_count)
            + "\n"
            + "Average LLM Token Count: "
            + str(token_counter.total_llm_token_count / len(cve_lst))
            + "\n"
        )

        token_counter.reset_counts()
    

    # epss_csv_path = 'epss_scores-2024-08-06.csv'
    # cve_analysis_from_epss_csv(epss_csv_path)

 

if __name__ == "__main__":
    main()

