Spaces:
Running
Running
File size: 2,656 Bytes
168ae1c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | 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() |