CWE.md

# Basic Data Validation Script for Large Language Models (LLMs)
# This script is designed to provide foundational checks for common vulnerabilities.

import re

# LLM01: Prompt Injection (CWE-77, CWE-94)
def validate_prompt(prompt):
    """
    Validates the given prompt to prevent injection attacks.
    """
    # Example validation to remove potentially dangerous characters or patterns
    clean_prompt = re.sub(r'[^\w\s]', '', prompt)
    return clean_prompt

# LLM02: Insecure Output Handling (CWE-79, CWE-116)
def encode_output(output):
    """
    Encodes the output to prevent XSS attacks or other output encoding issues.
    """
    # Simple HTML encoding example
    encoded_output = output.replace('<', '&lt;').replace('>', '&gt;')
    return encoded_output

# LLM03: Training Data Poisoning (CWE-506, CWE-915)
# Note: Validation should occur during the data collection and preparation phase.
def validate_training_data(data):
    """
    Validates training data to ensure it's not maliciously crafted.
    """
    # Example check for unexpected patterns or malicious content
    if "unexpected_pattern" in data:
        raise ValueError("Invalid training data detected.")
    return True

# LLM04: Model Denial of Service (CWE-400)
def check_query_limits(query):
    """
    Checks if the query exceeds certain limits to prevent DoS attacks.
    """
    MAX_LENGTH = 1000  # Example limit
    if len(query) > MAX_LENGTH:
        raise ValueError("Query exceeds maximum allowed length.")
    return True

# LLM05: Supply-Chain Vulnerabilities (CWE-829, CWE-937)
# Manual review and using trusted libraries are recommended.

# LLM06: Sensitive Information Disclosure (CWE-200)
def redact_sensitive_info(text):
    """
    Redacts sensitive information from the text.
    """
    # Example redaction (simple and should be customized)
    redacted_text = re.sub(r'\b(account_number|ssn)\b', '[REDACTED]', text, flags=re.IGNORECASE)
    return redacted_text

# LLM07: Insecure Plugin Design (CWE-749, CWE-1203)
# Ensure plugins do not expose dangerous methods or functions directly to end-users.

# LLM08: Excessive Agency (CWE-807)
# Implement strict checks on inputs used for security decisions.

# LLM09: Overreliance (CWE-1048)
# Ensure diversification in security mechanisms and checks.

# LLM10: Model Theft (CWE-494, CWE-1241)
# Protect model artifacts using integrity checks and secure distribution methods.

# Example usage
prompt = "Example <script>alert('XSS')</script> prompt"
clean_prompt = validate_prompt(prompt)
print(f"Cleaned Prompt: {clean_prompt}")

output = "<h1>This is a header</h1>"
encoded_output = encode_output(output)
print(f"Encoded Output: {encoded_output}")


### Recommended Automated Validation Tools

For enhancing the security of your LLM application, integrating automated validation tools into your CI/CD pipeline is crucial. Here are some tools that can be particularly useful:

1. **Bandit**: A tool designed to find common security issues in Python code. It's useful for static analysis and can help detect security issues related to the CWEs mentioned.
   
2. **OWASP ZAP (Zed Attack Proxy)**: An open-source web application security scanner. While more web-focused, it can be useful for testing web interfaces to LLM applications for issues like XSS (CWE-79) and other web vulnerabilities.

3. **SonarQube**: Offers comprehensive code quality and security scanning, including detection of vulnerabilities and code smells.

4. **CodeQL**: GitHub's code scanning tool that uses queries to identify vulnerabilities across multiple languages, including Python. It can be used to automate security checks as part of your GitHub Actions workflows.

5. **PyTorch/TensorFlow Security Advisories**: For LLMs built on these frameworks, staying updated with the latest security advisories is crucial. Though not tools per se, subscribing to these advisories can help mitigate supply-chain vulnerabilities (CWE-829, CWE-937).

Each of these tools can be integrated into your development and deployment processes to automatically flag potential security issues, helping adhere to secure coding practices and mitigate vulnerabilities associated with LLMs.