Spaces:
Running
Running
File size: 6,182 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 | import json
import re
from datasets import Dataset, DatasetDict
import ast
def check_syntax(code):
"""Check if Python code has syntax errors"""
try:
ast.parse(code)
return True, None
except SyntaxError as e:
return False, str(e)
def create_enhanced_dataset():
"""Create better dataset with syntax errors marked"""
samples = []
# ===== VULNERABLE CODE (with proper syntax) =====
vulnerable_examples = [
# SQL Injection (proper syntax)
{
"code": """def get_user(user_id):\n query = f"SELECT * FROM users WHERE id = {user_id}"\n return execute_query(query)""",
"label": 1,
"type": "sql_injection",
"explanation": "SQL injection: User input directly in query string"
},
{
"code": """cursor.execute(f"DELETE FROM logs WHERE date = '{user_date}'")""",
"label": 1,
"type": "sql_injection",
"explanation": "SQL injection in DELETE statement"
},
# Hardcoded Secrets
{
"code": """API_KEY = "EX_SAMPLE_KEY_NOT_REAL_12345" """,
"label": 1,
"type": "hardcoded_secret",
"explanation": "Hardcoded API key in source code"
},
{
"code": """db_password = "SuperSecret123!" """,
"label": 1,
"type": "hardcoded_secret",
"explanation": "Hardcoded database password"
},
# XSS Vulnerabilities
{
"code": """def render_comment(comment):\n return f"<div>{comment}</div>" """,
"label": 1,
"type": "xss",
"explanation": "XSS: User input in HTML without sanitization"
},
# Insecure Deserialization
{
"code": """import pickle\ndata = pickle.loads(user_input)""",
"label": 1,
"type": "insecure_deserialization",
"explanation": "Pickle deserialization of user input is dangerous"
},
# Command Injection
{
"code": """import os\nos.system(f"echo {user_input}")""",
"label": 1,
"type": "command_injection",
"explanation": "Command injection: User input in system command"
}
]
# ===== SAFE CODE =====
safe_examples = [
# Safe SQL
{
"code": """def get_user(user_id):\n query = "SELECT * FROM users WHERE id = %s"\n return execute_query(query, (user_id,))""",
"label": 0,
"type": "safe_sql",
"explanation": "Parameterized query prevents SQL injection"
},
# Safe Secrets
{
"code": """import os\nAPI_KEY = os.getenv("API_KEY")""",
"label": 0,
"type": "safe_secret",
"explanation": "API key from environment variable"
},
# Safe XSS
{
"code": """import html\ndef render_comment(comment):\n return f"<div>{html.escape(comment)}</div>" """,
"label": 0,
"type": "safe_xss",
"explanation": "HTML escaped user input prevents XSS"
},
# Safe Deserialization
{
"code": """import json\ndata = json.loads(user_input)""",
"label": 0,
"type": "safe_deserialization",
"explanation": "JSON deserialization is generally safe"
},
# Safe Command Execution
{
"code": """import subprocess\nsubprocess.run(["echo", "static_text"], check=True)""",
"label": 0,
"type": "safe_command",
"explanation": "Command with static arguments"
}
]
# ===== SYNTAX ERRORS (NEW!) =====
syntax_error_examples = [
{
"code": """def test()\n print("hello")""", # Missing colon
"label": 1, # Mark as vulnerable/problematic
"type": "syntax_error",
"explanation": "Syntax error: Missing colon after function definition"
},
{
"code": """x = [1, 2, 3""", # Missing closing bracket
"label": 1,
"type": "syntax_error",
"explanation": "Syntax error: Missing closing bracket"
},
{
"code": """print("hello'""", # Quote mismatch
"label": 1,
"type": "syntax_error",
"explanation": "Syntax error: Mismatched quotes"
}
]
# Combine all examples
all_examples = []
# Add multiple copies for balance
for _ in range(5): # 5 copies of each
all_examples.extend(vulnerable_examples)
all_examples.extend(safe_examples)
# Add syntax errors
all_examples.extend(syntax_error_examples)
# Check syntax of all examples
print("🔍 Checking syntax of all examples...")
for i, example in enumerate(all_examples):
is_valid, error = check_syntax(example["code"])
example["has_syntax_error"] = not is_valid
if error:
example["syntax_error"] = error
# Save to files
import os
os.makedirs("enhanced_data", exist_ok=True)
# Save as JSON
with open("enhanced_data/dataset.json", "w") as f:
json.dump(all_examples, f, indent=2)
# Create Hugging Face Dataset
dataset = Dataset.from_list(all_examples)
# Split into train/test
dataset = dataset.train_test_split(test_size=0.2, seed=42)
dataset.save_to_disk("enhanced_data/hf_dataset")
# Statistics
print("\n📊 DATASET STATISTICS:")
print(f"Total examples: {len(all_examples)}")
print(f"Vulnerable examples: {sum(1 for e in all_examples if e['label'] == 1)}")
print(f"Safe examples: {sum(1 for e in all_examples if e['label'] == 0)}")
print(f"With syntax errors: {sum(1 for e in all_examples if e.get('has_syntax_error', False))}")
print("\n✅ Enhanced dataset created!")
print("📁 Saved to: enhanced_data/")
return all_examples
if __name__ == "__main__":
create_enhanced_dataset() |