cve_dataset_llama.py

import os
import json

containers = None
cve_metadata = None
cve_id = None
description = None
references = None
cve_state = None

def extract_json_data(filename):
    global containers, cve_metadata, cve_id, description, references, cve_state

    with open(filename, "r", newline="", encoding="utf-8") as json_file:
        data = json.load(json_file)

    containers = data.get("containers", {})
    cve_metadata = data.get("cveMetadata", {})

    cna = containers.get("cna", {})
    description = cna.get("descriptions", [{}])[0].get("value", "N\\A")
    references = cna.get("references", "N\\A")
    cve_state = cve_metadata.get("state", "N\\A")
    cve_id = cve_metadata.get("cveId", "N\\A")

def standardize_response(response):
    """Ensure response is a string, flattening arrays or serializing dictionaries if necessary."""
    if isinstance(response, list):
        return ", ".join([str(item) for item in response])  # Convert list to a comma-separated string
    elif isinstance(response, dict):
        return json.dumps(response, ensure_ascii=False)  # Convert dict to a JSON string
    return str(response)  # Convert other types to string

def write_jsonl_entries(filename, entries):
    with open(filename, "a", encoding="utf-8") as jsonl_file:
        for entry in entries:
            # Standardize all `Analyst Response` fields before writing
            entry["Analyst Response"] = standardize_response(entry["Analyst Response"])
            json.dump(entry, jsonl_file, ensure_ascii=False)
            jsonl_file.write("\n")

def generate_training_entries():
    return [
        {
            "CVE ID": cve_id,
            "Human Input": f"Explain the CVE ID: What does the identifier {cve_id} mean in the context of cybersecurity?",
            "Analyst Response": cve_id
        },
        {
            "CVE ID": cve_id,
            "Human Input": f"Describe the vulnerability: Provide detailed information about the vulnerability for {cve_id}.",
            "Analyst Response": description
        },
        {
            "CVE ID": cve_id,
            "Human Input": f"Provide references for further information: Can you share references or external links related to {cve_id}?",
            "Analyst Response": references
        },
        {
            "CVE ID": cve_id,
            "Human Input": f"Explain the state of this CVE: What is the publication or resolution state of {cve_id}?",
            "Analyst Response": cve_state
        }
    ]

def main():
    jsonl_filename = "train.jsonl"
    if os.path.exists(jsonl_filename):
        os.remove(jsonl_filename)  # Ensure no duplicate entries

    for year in range(1999, 2024):
        year_folder = f"cves/{year}"
        if os.path.exists(year_folder) and os.path.isdir(year_folder):
            for subfolder in ["0xxx", "1xxx"]:
                subfolder_path = os.path.join(year_folder, subfolder)
                if os.path.exists(subfolder_path) and os.path.isdir(subfolder_path):
                    for root, _, files in os.walk(subfolder_path):
                        for file in files:
                            if file.startswith(f"CVE-{year}-") and file.endswith(".json"):
                                json_file = os.path.join(root, file)
                                extract_json_data(json_file)

                                # Generate individual entries for robust training
                                entries = generate_training_entries()
                                write_jsonl_entries(jsonl_filename, entries)

if __name__ == "__main__":
    main()