File size: 3,248 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 | import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from peft import PeftModel
BASE_MODEL = "codellama/CodeLlama-7b-Instruct-hf"
ADAPTER_PATH = "/home/ai/codellama-security-finetuned"
def load_model():
print("Loading base model...")
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(
BASE_MODEL,
quantization_config=bnb_config,
device_map="auto",
trust_remote_code=True,
torch_dtype=torch.bfloat16,
)
print("Loading LoRA adapter...")
model = PeftModel.from_pretrained(model, ADAPTER_PATH)
model = model.merge_and_unload()
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
return model, tokenizer
def analyze_code(model, tokenizer, code, instruction="Identify all security vulnerabilities in this code"):
prompt = 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|>
{instruction}
```python
{code}
```
<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=1024,
temperature=0.3,
top_p=0.9,
do_sample=True,
repetition_penalty=1.1,
pad_token_id=tokenizer.eos_token_id,
)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return response.split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()
def main():
model, tokenizer = load_model()
test_cases = [
("SQL Injection", """
import sqlite3
def get_user(username):
conn = sqlite3.connect('app.db')
cursor = conn.cursor()
query = f"SELECT * FROM users WHERE name = '{username}'"
return cursor.execute(query).fetchall()
"""),
("Command Injection", """
import subprocess
def ping_host(host):
result = subprocess.run(f"ping -c 4 {host}", shell=True, capture_output=True)
return result.stdout
"""),
("Path Traversal", """
from flask import request, send_file
@app.route('/download')
def download():
filename = request.args.get('file')
return send_file(f'/var/www/uploads/{filename}')
"""),
("Insecure Deserialization", """
import pickle
def load_session(data):
return pickle.loads(data)
"""),
("JWT Issues", """
import jwt
def verify_token(token):
return jwt.decode(token, 'secret123', algorithms=['HS256'])
"""),
]
for name, code in test_cases:
print(f"\n{'='*60}")
print(f"TEST: {name}")
print(f"{'='*60}")
result = analyze_code(model, tokenizer, code)
print(result)
if __name__ == "__main__":
main() |