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()