File size: 11,673 Bytes
f238825 | 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 | import os
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import json
import torch
from datasets import Dataset
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TrainingArguments,
Trainer,
DataCollatorForLanguageModeling,
BitsAndBytesConfig
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, TaskType
MODEL_NAME = "codellama/CodeLlama-7b-Instruct-hf"
OUTPUT_DIR = "/home/ai/codellama-security-finetuned"
SECURITY_DATA = [
{"instruction": "Identify the vulnerability in this code", "input": "```python\nimport sqlite3\nconn = sqlite3.connect('users.db')\ncursor = conn.cursor()\nusername = input('Enter username: ')\nquery = f\"SELECT * FROM users WHERE username = '{username}'\"\ncursor.execute(query)\n```", "output": "SQL Injection vulnerability. The query uses string interpolation directly with user input. Fix: Use parameterized queries: `cursor.execute(\"SELECT * FROM users WHERE username = ?\", (username,))`"},
{"instruction": "Find the security flaw", "input": "```python\nimport subprocess\nfilename = input('Enter filename: ')\nsubprocess.run(f'cat {filename}', shell=True)\n```", "output": "Command Injection via shell=True and unsanitized user input. Fix: Use `subprocess.run(['cat', filename], shell=False)` with input validation."},
{"instruction": "What vulnerability exists here?", "input": "```python\nimport pickle\ndata = pickle.loads(request.data)\n```", "output": "Insecure Deserialization. Pickle can execute arbitrary code during deserialization. Fix: Use JSON or safe serialization formats, validate input."},
{"instruction": "Identify the XSS vulnerability", "input": "```javascript\napp.get('/search', (req, res) => {\n res.send(`<h1>Results for: ${req.query.q}</h1>`);\n});\n```", "output": "Reflected XSS. User input directly rendered in HTML without escaping. Fix: Use template engines with auto-escaping or sanitize input with DOMPurify."},
{"instruction": "Find the authentication bypass", "input": "```python\n@app.route('/admin')\ndef admin():\n if request.headers.get('X-User-Role') == 'admin':\n return 'Admin panel'\n return 'Forbidden', 403\n```", "output": "Client-side authorization check. Header can be spoofed. Fix: Use server-side session validation with proper authentication tokens."},
{"instruction": "What's wrong with this crypto implementation?", "input": "```python\nfrom Crypto.Cipher import AES\nkey = b'mysecretkey12345'\ncipher = AES.new(key, AES.MODE_ECB)\nciphertext = cipher.encrypt(pad(data, 16))\n```", "output": "AES-ECB mode is insecure (deterministic, reveals patterns). Hardcoded key. Fix: Use AES-GCM with random IV, derive key via PBKDF2/Argon2."},
{"instruction": "Identify the path traversal", "input": "```python\n@app.route('/download')\ndef download():\n filename = request.args.get('file')\n return send_file(f'/var/www/files/{filename}')\n```", "output": "Path Traversal. No validation on filename allows `../../etc/passwd`. Fix: Use `secure_filename`, validate against allowlist, use `os.path.basename`."},
{"instruction": "Find the SSRF vulnerability", "input": "```python\nimport requests\nurl = request.args.get('url')\nresponse = requests.get(url)\nreturn response.text\n```", "output": "Server-Side Request Forgery (SSRF). User controls URL parameter, can access internal services. Fix: Validate URL against allowlist, block private IPs, use allowlist of domains."},
{"instruction": "What's the race condition here?", "input": "```python\ndef transfer(from_acc, to_acc, amount):\n if get_balance(from_acc) >= amount:\n debit(from_acc, amount)\n credit(to_acc, amount)\n```", "output": "TOCTOU Race Condition. Balance check and transfer aren't atomic. Fix: Use database transactions with SELECT FOR UPDATE or atomic operations."},
{"instruction": "Identify the JWT vulnerability", "input": "```python\nimport jwt\ntoken = request.headers.get('Authorization')\npayload = jwt.decode(token, 'secret', algorithms=['HS256'])\n```", "output": "Hardcoded secret, no algorithm validation (allows 'none' alg), no expiration check. Fix: Use strong secret from env, specify algorithms=['HS256'], verify exp claim."},
{"instruction": "Find the XXE vulnerability", "input": "```python\nimport xml.etree.ElementTree as ET\ndef parse_xml(xml_data):\n return ET.fromstring(xml_data)\n```", "output": "XML External Entity (XXE) injection. Parser processes external entities. Fix: Disable DTD processing: `ET.XMLParser(resolve_entities=False)` or use defusedxml."},
{"instruction": "What's the issue with this CORS config?", "input": "```python\napp.add_middleware(CORSMiddleware,\n allow_origins=['*'],\n allow_credentials=True,\n allow_methods=['*'],\n allow_headers=['*']\n)\n```", "output": "Wildcard origin with credentials enabled - violates CORS spec, allows credential theft. Fix: Specify exact origins, don't use '*' with credentials."},
{"instruction": "Identify the insecure direct object reference", "input": "```python\n@app.route('/api/user/<int:user_id>/profile')\ndef get_profile(user_id):\n return jsonify(db.query(User).filter_by(id=user_id).first())\n```", "output": "IDOR - No authorization check. Users can access any profile by changing user_id. Fix: Verify current user owns the resource or has permission."},
{"instruction": "Find the weak randomness", "input": "```python\nimport random\nsession_token = ''.join(random.choices('abcdefghijklmnopqrstuvwxyz0123456789', k=32))\n```", "output": "Using random (Mersenne Twister) for cryptographic purposes - predictable. Fix: Use `secrets.token_urlsafe(32)` or `os.urandom()` for crypto-safe randomness."},
{"instruction": "What's the container escape risk?", "input": "```dockerfile\nFROM ubuntu:latest\nRUN apt-get update && apt-get install -y docker.io\nVOLUME /var/run/docker.sock\n```", "output": "Mounting docker.sock gives container root access to host. Fix: Don't mount docker.sock; use socket proxies like docker-socket-proxy with limited permissions."},
{"instruction": "Identify the supply chain attack vector", "input": "```json\n{\n \"dependencies\": {\n \"left-pad\": \"^1.0.0\",\n \"event-stream\": \"^3.3.4\"\n }\n}\n```", "output": "Unpinned dependencies allow malicious updates. event-stream was compromised. Fix: Pin exact versions, use lockfiles, enable Dependabot, verify package integrity with npm audit/signatures."},
{"instruction": "Find the timing attack vulnerability", "input": "```python\ndef verify_token(token, expected):\n return token == expected\n```", "output": "String comparison short-circuits - timing attack reveals token character by character. Fix: Use `hmac.compare_digest(token, expected)` for constant-time comparison."},
{"instruction": "What's the log injection issue?", "input": "```python\nimport logging\nlogging.info(f'User {username} logged in from {ip}')\n```", "output": "Log Injection - username/ip can contain newlines to forge log entries. Fix: Sanitize input, use structured logging (JSON), escape newlines."},
{"instruction": "Identify the broken access control", "input": "```python\n@app.route('/api/admin/delete-user/<int:id>')\ndef delete_user(id):\n User.query.filter_by(id=id).delete()\n return 'OK'\n```", "output": "Missing authorization - any authenticated user can delete users. Fix: Add @require_admin decorator, check user permissions before action."},
{"instruction": "Find the prototype pollution", "input": "```javascript\nfunction merge(target, source) {\n for (let key in source) {\n target[key] = source[key];\n }\n}\nmerge({}, JSON.parse(userInput));\n```", "output": "Prototype Pollution - `__proto__` or `constructor.prototype` in input pollutes Object.prototype. Fix: Validate keys, use `Object.create(null)`, freeze prototype, use lodash.mergeWith."},
{"instruction": "What's the deserialization RCE?", "input": "```java\nObjectInputStream ois = new ObjectInputStream(new FileInputStream(\"data.ser\"));\nObject obj = ois.readObject();\n```", "output": "Java Deserialization RCE - gadget chains (CommonsCollections, etc.) execute code. Fix: Avoid native serialization, use JSON, implement ObjectInputFilter, use SerialKiller/NotSoSerial."}
]
def format_example(example):
return f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
You are a cybersecurity expert specializing in vulnerability assessment, penetration testing, and secure code review. Identify security flaws, explain the impact, and provide remediation.
<|eot_id|><|start_header_id|>user<|end_header_id|>
{example['instruction']}
{example['input']}
<|eot_id|><|start_header_id|>assistant<|end_header_id|>
{example['output']}
<|eot_id|>"""
def main():
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"
print("Loading model with 4-bit quantization...")
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
torch_dtype=torch.bfloat16,
)
model = prepare_model_for_kbit_training(model)
lora_config = LoraConfig(
r=64,
lora_alpha=128,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
lora_dropout=0.05,
bias="none",
task_type=TaskType.CAUSAL_LM,
)
model = get_peft_model(model, lora_config)
model.print_trainable_parameters()
print("Preparing dataset...")
formatted_data = [{"text": format_example(ex)} for ex in SECURITY_DATA]
dataset = Dataset.from_list(formatted_data)
def tokenize_function(examples):
result = tokenizer(
examples["text"],
truncation=True,
max_length=1024,
padding="max_length",
return_tensors="pt"
)
result["labels"] = result["input_ids"].clone()
return result
tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
tokenized_dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "labels"])
training_args = TrainingArguments(
output_dir=OUTPUT_DIR,
num_train_epochs=3,
per_device_train_batch_size=1,
gradient_accumulation_steps=8,
warmup_steps=10,
learning_rate=2e-4,
fp16=False,
bf16=True,
logging_steps=5,
save_steps=50,
save_total_limit=2,
optim="paged_adamw_8bit",
lr_scheduler_type="cosine",
report_to="none",
remove_unused_columns=False,
dataloader_pin_memory=False,
max_grad_norm=0.3,
)
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
trainer = Trainer(
model=model,
train_dataset=tokenized_dataset,
args=training_args,
data_collator=data_collator,
)
print("Starting training...")
trainer.train()
print("Saving model...")
trainer.save_model()
tokenizer.save_pretrained(OUTPUT_DIR)
print(f"Model saved to {OUTPUT_DIR}")
if __name__ == "__main__":
main() |