create_dataset.py

import json
import os
from datasets import Dataset

def create_simple_dataset():
    """Create a small dataset for quick testing"""
    
    samples = []
    
    # Vulnerable code examples (label: 1)
    vulnerable = [
        {
            "code": "query = f'SELECT * FROM users WHERE id = {user_id}'",
            "label": 1,
            "type": "sql_injection",
            "explanation": "SQL injection vulnerability: user input directly in query"
        },
        {
            "code": "api_key = 'sk_live_1234567890abcdef'",
            "label": 1,
            "type": "hardcoded_secret",
            "explanation": "Hardcoded API key in source code"
        },
        {
            "code": "result = eval(user_input)",
            "label": 1,
            "type": "insecure_deserialization",
            "explanation": "eval() with user input is dangerous"
        },
        {
            "code": "return f'<div>{user_input}</div>'",
            "label": 1,
            "type": "xss",
            "explanation": "Potential XSS: user input in HTML without sanitization"
        },
    ]
    
    # Safe code examples (label: 0)
    safe = [
        {
            "code": "query = 'SELECT * FROM users WHERE id = %s'",
            "label": 0,
            "type": "safe",
            "explanation": "Parameterized query prevents SQL injection"
        },
        {
            "code": "api_key = os.getenv('API_KEY')",
            "label": 0,
            "type": "safe",
            "explanation": "API key from environment variable"
        },
        {
            "code": "result = json.loads(user_input)",
            "label": 0,
            "type": "safe",
            "explanation": "Safe deserialization with json.loads"
        },
        {
            "code": "return f'<div>{html.escape(user_input)}</div>'",
            "label": 0,
            "type": "safe",
            "explanation": "HTML escaped user input prevents XSS"
        },
    ]
    
    # Create more examples by modifying
    for i in range(20):
        samples.extend(vulnerable)
        samples.extend(safe)
    
    # Create directory if not exists
    os.makedirs("data", exist_ok=True)
    
    # Save as JSON
    with open("data/dataset.json", "w") as f:
        json.dump(samples, f, indent=2)
    
    # Also create Hugging Face Dataset
    dataset = Dataset.from_list(samples)
    dataset.save_to_disk("data/hf_dataset")
    
    print(f"✅ Created dataset with {len(samples)} examples")
    print(f"📁 Saved to: data/dataset.json")
    print(f"📁 HF Dataset: data/hf_dataset")
    
    return samples

if __name__ == "__main__":
    create_simple_dataset()