detection

Vulnerability Detection Engine for LLM Responses

This package provides a comprehensive vulnerability detection engine for analyzing LLM responses to identify security vulnerabilities based on template-defined criteria.

Features

• Multiple Detection Methods:

  • String matching for exact or partial text matches
  • Regular expression pattern matching for complex text patterns
  • Semantic analysis for intent or meaning-based vulnerabilities
  • Hybrid detection combining multiple methods

• Vulnerability Classification:

  • Mapping to OWASP Top 10 for LLMs categories
  • Correlation with ISO/IEC 42001 controls
  • Severity ratings (Critical, High, Medium, Low, Info)
  • Confidence scores for detections

• Extensibility:

  • Plugin architecture for custom detection methods
  • Hooks for pre and post-detection processing
  • Support for custom embedding providers

• Performance Optimizations:

  • Concurrent detection for multiple criteria
  • Caching of embeddings for semantic matching
  • Batch processing capabilities

Architecture

The vulnerability detection engine is designed with a modular architecture:

1. Core Types (types.go):

   • Defines the interfaces and data structures for the detection engine
   • Includes vulnerability types, severity levels, and detection methods

2. Detection Engine (engine.go):

   • Implements the main detection engine
   • Manages detectors, hooks, and providers
   • Handles concurrent detection and result aggregation

3. Detectors (detectors.go):

   • Implements specific detection methods
   • Includes string matching, regex matching, semantic matching, and hybrid matching

4. Embedding Service (embedding.go):

   • Provides semantic analysis capabilities
   • Supports multiple embedding providers
   • Includes caching for performance optimization

5. Providers (providers.go):

   • Implements standard mapping and remediation providers
   • Maps vulnerabilities to industry standards
   • Provides remediation suggestions

6. Factory (factory.go):

   • Creates and configures detection engine components
   • Simplifies the creation of complex detection pipelines

Usage

Basic Usage

package main

import (
    "context"
    "fmt"

    "github.com/perplext/LLMrecon/src/vulnerability/detection"
)

func main() {
    // Create factory
    factory := detection.NewFactory()

    // Create detection engine
    engine := factory.CreateDetectionEngine()

    // Define detection criteria
    criteria := []detection.DetectionCriteria{
        {
            Type:  detection.StringMatch,
            Match: "password",
            Context: map[string]interface{}{
                "vulnerability_type": string(detection.SensitiveInfoDisclosure),
                "severity":           string(detection.High),
            },
        },
        {
            Type:    detection.RegexMatch,
            Pattern: `\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b`,
            Context: map[string]interface{}{
                "vulnerability_type": string(detection.SensitiveInfoDisclosure),
                "severity":           string(detection.Medium),
            },
        },
    }

    // LLM response to analyze
    response := "My password is 12345 and my email is test@example.com"

    // Detect vulnerabilities
    results, err := engine.Detect(context.Background(), response, criteria)
    if err != nil {
        fmt.Printf("Error: %v\n", err)
        return
    }

    // Process results
    for _, result := range results {
        if result.Detected {
            fmt.Printf("Detected %s vulnerability with %s severity\n", 
                result.VulnerabilityType, result.Severity)
            fmt.Printf("Confidence: %d%%\n", result.ConfidenceScore)
            fmt.Printf("Matched content: %v\n", result.MatchedContent)
            
            if len(result.StandardMappings) > 0 {
                fmt.Printf("Standard mappings:\n")
                for _, mapping := range result.StandardMappings {
                    fmt.Printf("  - %s %s: %s\n", 
                        mapping.Standard, mapping.ID, mapping.Name)
                }
            }
            
            if result.RemediationSuggestion != "" {
                fmt.Printf("Remediation: %s\n", result.RemediationSuggestion)
            }
            
            fmt.Println()
        }
    }
}

Custom Detector

package main

import (
    "context"
    "fmt"
    "time"

    "github.com/perplext/LLMrecon/src/vulnerability/detection"
)

// CustomDetector is a custom vulnerability detector
type CustomDetector struct{}

// Name returns the name of the detector
func (d *CustomDetector) Name() string {
    return "custom_detector"
}

// Detect detects vulnerabilities in an LLM response
func (d *CustomDetector) Detect(ctx context.Context, response string, criteria detection.DetectionCriteria) (*detection.DetectionResult, error) {
    // Custom detection logic
    detected := false
    
    // Example: detect if response is too long
    if len(response) > 1000 {
        detected = true
    }
    
    // Create result
    result := &detection.DetectionResult{
        Detected:          detected,
        VulnerabilityType: detection.Overreliance,
        Severity:          detection.Medium,
        ConfidenceScore:   detected ? 100 : 0,
        DetectionMethod:   detection.DetectionMethod("custom"),
        Context: map[string]interface{}{
            "response_length": len(response),
        },
        Timestamp: time.Now(),
    }
    
    return result, nil
}

func main() {
    // Create detection engine
    engine := detection.NewDetectionEngine()
    
    // Register custom detector
    engine.RegisterDetector(&CustomDetector{})
    
    // Define detection criteria
    criteria := []detection.DetectionCriteria{
        {
            Type: detection.DetectionMethod("custom"),
        },
    }
    
    // LLM response to analyze
    response := "This is a very long response..." // Imagine this is >1000 chars
    
    // Detect vulnerabilities
    results, err := engine.Detect(context.Background(), response, criteria)
    if err != nil {
        fmt.Printf("Error: %v\n", err)
        return
    }
    
    // Process results
    for _, result := range results {
        if result.Detected {
            fmt.Printf("Detected %s vulnerability with %s severity\n", 
                result.VulnerabilityType, result.Severity)
        }
    }
}

Batch Detection

package main

import (
    "context"
    "fmt"

    "github.com/perplext/LLMrecon/src/vulnerability/detection"
)

func main() {
    // Create factory
    factory := detection.NewFactory()

    // Create detection engine
    engine := factory.CreateDetectionEngine()

    // Define batch detection requests
    requests := []detection.BatchDetectionRequest{
        {
            Response: "My password is 12345",
            Criteria: []detection.DetectionCriteria{
                {
                    Type:  detection.StringMatch,
                    Match: "password",
                    Context: map[string]interface{}{
                        "vulnerability_type": string(detection.SensitiveInfoDisclosure),
                        "severity":           string(detection.High),
                    },
                },
            },
        },
        {
            Response: "My email is test@example.com",
            Criteria: []detection.DetectionCriteria{
                {
                    Type:    detection.RegexMatch,
                    Pattern: `\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b`,
                    Context: map[string]interface{}{
                        "vulnerability_type": string(detection.SensitiveInfoDisclosure),
                        "severity":           string(detection.Medium),
                    },
                },
            },
        },
    }

    // Perform batch detection
    results, err := engine.BatchDetect(context.Background(), requests)
    if err != nil {
        fmt.Printf("Error: %v\n", err)
        return
    }

    // Process results
    for i, result := range results {
        fmt.Printf("Response %d:\n", i+1)
        
        if result.Error != "" {
            fmt.Printf("  Error: %s\n", result.Error)
            continue
        }
        
        for j, detection := range result.Results {
            if detection.Detected {
                fmt.Printf("  Detection %d: %s (%s)\n", 
                    j+1, detection.VulnerabilityType, detection.Severity)
            }
        }
        
        fmt.Println()
    }
}

Command Line Interface

The vulnerability detection engine can be used from the command line using the detect command:

# Detect vulnerabilities in an LLM response
$ LLMrecon detect -r response.txt -c criteria.json -o results.json

# Enable verbose output
$ LLMrecon detect -r response.txt -c criteria.json -v

Example Criteria File

[
  {
    "type": "string_match",
    "match": "password",
    "case_sensitive": false,
    "context": {
      "vulnerability_type": "sensitive_info_disclosure",
      "severity": "high"
    }
  },
  {
    "type": "regex_match",
    "pattern": "\\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\\.[A-Za-z]{2,}\\b",
    "context": {
      "vulnerability_type": "sensitive_info_disclosure",
      "severity": "medium"
    }
  }
]

Integration with Template Management

The vulnerability detection engine is designed to integrate seamlessly with the Template Management System:

package main

import (
    "context"
    "fmt"

    "github.com/perplext/LLMrecon/src/template/format"
    "github.com/perplext/LLMrecon/src/template/management"
    "github.com/perplext/LLMrecon/src/vulnerability/detection"
)

func main() {
    // Create detection engine
    factory := detection.NewFactory()
    engine := factory.CreateDetectionEngine()

    // Load template
    template, err := format.LoadFromFile("path/to/template.yaml")
    if err != nil {
        fmt.Printf("Error loading template: %v\n", err)
        return
    }

    // Convert template detection criteria to engine criteria
    criteria := []detection.DetectionCriteria{
        {
            Type:      detection.DetectionMethod(template.Test.Detection.Type),
            Pattern:   template.Test.Detection.Pattern,
            Match:     template.Test.Detection.Match,
            Criteria:  template.Test.Detection.Criteria,
            Condition: template.Test.Detection.Condition,
            Context: map[string]interface{}{
                "vulnerability_type": template.Info.Compliance.OWASP,
                "severity":           template.Info.Severity,
            },
        },
    }

    // LLM response to analyze
    response := "This is the LLM response to analyze"

    // Detect vulnerabilities
    results, err := engine.Detect(context.Background(), response, criteria)
    if err != nil {
        fmt.Printf("Error: %v\n", err)
        return
    }

    // Create template result
    templateResult := &management.TemplateResult{
        TemplateID: template.ID,
        Status:     management.StatusCompleted,
        Response:   response,
    }

    // Update template result with detection results
    if len(results) > 0 && results[0].Detected {
        templateResult.Detected = true
        templateResult.Score = results[0].ConfidenceScore
        templateResult.Details = results[0].Context
    }

    // Process template result
    fmt.Printf("Template %s: %v\n", templateResult.TemplateID, templateResult.Detected)
}

Performance Considerations

The vulnerability detection engine is designed with performance in mind:

• Concurrent Detection: Multiple detection criteria are processed concurrently
• Embedding Caching: Embeddings for semantic matching are cached to avoid redundant computation
• Batch Processing: Multiple responses can be processed in a single batch operation
• Efficient Matchers: String and regex matchers are optimized for performance

Extending the Engine

The vulnerability detection engine is designed to be extensible:

1. Custom Detectors: Implement the VulnerabilityDetector interface
2. Custom Embedding Providers: Implement the EmbeddingProvider interface
3. Custom Hooks: Create functions that implement the DetectionHook type
4. Custom Providers: Implement the RemediationProvider or StandardMappingProvider interfaces

Future Enhancements

• Persistent Storage: Add support for storing detection results in a database
• Distributed Detection: Support for distributed detection across multiple nodes
• Advanced Semantic Analysis: Integration with more sophisticated embedding models
• Real-time Detection: Support for real-time detection in streaming scenarios
• Visualization: Tools for visualizing detection results and trends