app.py

import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
import spaces

# Load the model and tokenizer
peft_model_id = "rootxhacker/CodeAstra-7B"
config = PeftConfig.from_pretrained(peft_model_id)

# Function to move tensors to CPU
def to_cpu(obj):
    if isinstance(obj, torch.Tensor):
        return obj.cpu()
    elif isinstance(obj, list):
        return [to_cpu(item) for item in obj]
    elif isinstance(obj, tuple):
        return tuple(to_cpu(item) for item in obj)
    elif isinstance(obj, dict):
        return {key: to_cpu(value) for key, value in obj.items()}
    return obj

# Load the model
model = AutoModelForCausalLM.from_pretrained(
    config.base_model_name_or_path,
    return_dict=True,
    load_in_4bit=True,
    device_map='auto'
)
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)

# Load the Lora model
model = PeftModel.from_pretrained(model, peft_model_id)

@spaces.GPU()
def get_completion(query, model, tokenizer):
    try:
        # Move model to CUDA
        model = model.cuda()
        # Ensure input is on CUDA
        inputs = tokenizer(query, return_tensors="pt").to('cuda')
        with torch.no_grad():
            outputs = model.generate(**inputs, max_new_tokens=1024, do_sample=True, temperature=0.7)
        # Move outputs to CPU before decoding
        outputs = to_cpu(outputs)
        return tokenizer.decode(outputs[0], skip_special_tokens=True)
    except Exception as e:
        return f"An error occurred: {str(e)}"
    finally:
        # Move model back to CPU to free up GPU memory
        model = model.cpu()
        torch.cuda.empty_cache()

@spaces.GPU()
def code_review(code_to_analyze):
    few_shot_prompt = """
You are an expert in secure coding practices and software logic analysis, tasked with reviewing source code for potential security vulnerabilities and logic flaws. Your goal is to understand the code, identify security issues that could be exploited, and uncover any logic vulnerabilities that could lead to unintended behavior. Follow these steps for each code review:

1. Understand the code
   - Analyze the purpose and functionality of the code
   - Identify input sources and output destinations
   - Note any security-sensitive operations (e.g., authentication, data storage, network communication)
   - Understand the logical flow and decision points in the code

2. Identify potential security issues
   - Look for common vulnerabilities (e.g., injection flaws, broken authentication, sensitive data exposure)
   - Consider how the code might be misused or exploited
   - Evaluate the handling of user input and data validation
   - Check for proper use of security functions and libraries

3. Identify potential logic vulnerabilities
   - Look for incorrect boolean logic in conditional statements
   - Check for off-by-one errors in loops and array operations
   - Identify potential race conditions in multi-threaded or asynchronous code
   - Evaluate edge cases and boundary conditions
   - Check for proper error handling and exception management

Here's an example of how to review code:

Code to review:
```php
function authenticateUser($username, $password) {
    $conn = new mysqli("localhost", "user", "password", "database");
    $query = "SELECT * FROM users WHERE username = '$username' AND password = '$password'";
    $result = $conn->query($query);
    if ($result->num_rows > 0) {
        return true;
    }
    return false;
}
```

Review:
1. Understanding the code:
   - The function attempts to authenticate a user based on a username and password.
   - It queries a database to check if the credentials exist.
   - Returns true if a matching user is found, false otherwise.

2. Potential security issues:
   - SQL Injection vulnerability: Username and password are directly inserted into the query.
   - Passwords are stored and compared in plain text, which is a severe security risk.
   - Hardcoded database credentials in the source code.
   - Potential for timing attacks due to direct string comparison.

3. Potential logic vulnerabilities:
   - The function returns true if more than one row is returned, which could lead to authentication bypass if multiple users have the same credentials.
   - No input validation on username or password, potentially allowing empty strings or null values.

Suggestions:
- Use prepared statements to prevent SQL injection.
- Use password hashing instead of storing plain text passwords.
- Store database credentials securely outside the source code.
- Implement proper error handling and use constant-time comparison for passwords.
- Ensure only one user can be authenticated at a time.
- Add input validation for username and password.

Now, review the following code using this approach:

{code_to_analyze}

Provide a detailed review following the structure above, including understanding the code, identifying potential security issues, identifying potential logic vulnerabilities, and offering specific suggestions for improvement. Start your response with 'Code Review:' and end it with 'End of Review.'"
"""

    query = few_shot_prompt.format(code_to_analyze=code_to_analyze)
    full_result = get_completion(query, model, tokenizer)
    
    # Extract only the AI's answer
    start_index = full_result.find("Code Review:")
    end_index = full_result.find("End of Review.")
    
    if start_index != -1 and end_index != -1:
        result = full_result[start_index:end_index].strip()
    else:
        result = "Unable to generate a proper code review. Please try again."
    
    return result

# Create Gradio interface
iface = gr.Interface(
    fn=code_review,
    inputs=gr.Textbox(lines=10, label="Enter code to analyze"),
    outputs=gr.Textbox(label="Code Review Result"),
    title="Code Review Expert",
    description="This tool analyzes code for potential security flaws, logic vulnerabilities, and provides guidance on secure coding practices."
)

# Launch the Gradio app
iface.launch()