cve_dataset_openai.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 write_openai_jsonl_entries(filename, entries):
    with open(filename, "a", encoding="utf-8") as jsonl_file:
        for entry in entries:
            json.dump(entry, jsonl_file, ensure_ascii=False)
            jsonl_file.write("\n")

def generate_openai_training_entries():
    return [
        {
            "messages": [
                {"role": "system", "content": "CVE Vulnerability Information"},
                {"role": "user", "content": f"Explain the CVE ID: What does the identifier {cve_id} mean in the context of cybersecurity?"},
                {"role": "assistant", "content": cve_id}
            ]
        },
        {
            "messages": [
                {"role": "system", "content": "CVE Vulnerability Information"},
                {"role": "user", "content": f"Describe the vulnerability: Provide detailed information about the vulnerability for {cve_id}."},
                {"role": "assistant", "content": description}
            ]
        },
        {
            "messages": [
                {"role": "system", "content": "CVE Vulnerability Information"},
                {"role": "user", "content": f"Provide references for further information: Can you share references or external links related to {cve_id}?"},
                {"role": "assistant", "content": references}
            ]
        },
        {
            "messages": [
                {"role": "system", "content": "CVE Vulnerability Information"},
                {"role": "user", "content": f"Explain the state of this CVE: What is the publication or resolution state of {cve_id}?"},
                {"role": "assistant", "content": cve_state}
            ]
        }
    ]

def main():
    jsonl_filename = "train_openai.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 OpenAI-format entries
                                entries = generate_openai_training_entries()
                                write_openai_jsonl_entries(jsonl_filename, entries)

if __name__ == "__main__":
    main()