Building a Custom AST-Grep Fixture Generator for Multi-Layered Vulnerability Scanning

Introduction

In modern codebases, spotting potential security flaws isn't always straightforward. Static analysis tools can catch common patterns, but deeper logic issues may slip by. That's why we've combined deterministic scanning through AST-Grep with a second layer of LLM-driven analysis, plus a unique self-evaluation approach. The result is a more reliable, multi-layered detection strategy.

Why a Multi-Layered Approach?

Relying on just one technique—say, scanning code for banned functions—can miss context-specific vulnerabilities. By combining an AST-Grep layer for deterministic checks with an LLM deep scan for contextual analysis, you get comprehensive coverage. Then, a self-verification or reflection step can reduce false positives and highlight low-confidence findings for manual review.

Setting Up AST-Grep and the LLM Orchestrator

Let's outline our high-level structure:

Copytrustsight/
├── ast-grep-service/
├── llm-service/
├── orchestrator/
└── docs/

We have three main parts:

• AST-Grep Service: Runs deterministic checks and flags known-bad patterns (like hardcoded secrets or forgotten TODOs).
• LLM Analyzer Service: Performs deeper, context-aware scanning on flagged snippets, using multiple AI endpoints if available.
• Orchestrator: Combines the results, manages self-verification, and generates a final consolidated report.

Implementing the AST-Grep Fixture Generator

Our AST-Grep fixture generator automatically produces test fixtures for scanning. This ensures coverage of specific rules in different contexts (e.g., JavaScript vs. Python). For example, a simple fixture might look like:

Copy# insecure_logging_example.py
try:
    result = 10 / 0
except Exception as e:
    print(str(e))  # Potentially insecure error logging

The AST-Grep service parses these fixtures, flags the suspicious print(str(e)) usage, then exports a structured JSON summary with rule IDs, severities, and file-line references.

LLM-Driven Analysis & Self-Evaluation

When a snippet is flagged, the LLM layer provides context-aware insights. For instance, it might suggest logging with a more secure mechanism or highlight that full stack traces can disclose sensitive information. After the LLM's first pass, an optional reflection prompt self-check can catch contradictory or inconsistent reasoning. The orchestrator merges all outcomes, weighs the confidence scores, and marks anything uncertain for manual review.

Implementation of the Go Test Generator

At the core of our fixture generator is a Go-based service that:

• Loads AST rule definitions from YAML or JSON.
• Handles existing test cases by migrating them into a consistent fixtures folder.
• Leverages a DeepSeek AI client to produce new “vulnerable” and “safe” code samples for each rule.
• Automatically determines file extensions (e.g., .py for Python) and splits the generated content into two separate test files.

For instance, our Generator struct in generator.go initializes the config and the DeepSeek client:

Copytype Generator struct {
    config *config.Config
    client *deepseek.Client
}

func NewGenerator() (*Generator, error) {
    cfg, err := config.LoadConfig()
    if err != nil {
        return nil, fmt.Errorf("failed to load config: %w", err)
    }

    client := deepseek.NewClient(cfg.DeepSeekAPIKey, 
        deepseek.WithModel(cfg.DeepSeekModel),
        deepseek.WithHTTPClient(&http.Client{
            Timeout: cfg.RequestTimeout,
        }),
    )
    
    return &Generator{
        config: cfg,
        client: client,
    }, nil
}

We then generate prompts and send them to DeepSeek to produce two distinct code snippets: one that should trigger a rule (vulnerable case) and one that should not (safe case). When the API response arrives, the code is parsed using simple regex for ``` blocks, and each snippet is saved to a fixtures directory based on the identified language:

Copyfunc (g *Generator) parseGeneratedCode(content string) []TestCase {
    codeBlockRegex := regexp.MustCompile("(?s)```(?:javascript|python|java|php)?\\n(.*?)\\n```")
    matches := codeBlockRegex.FindAllStringSubmatch(content, -1)
    // ...
    // Detect language, split into vulnerable vs safe, write .js/.py/.php etc.
    // ...
    return testCases
}

Our test generator also migrates legacy test cases from older directories into a single fixtures folder. This keeps everything consistent and ready for AST-Grep to parse. By the time our final LLM deep-scan starts, we have a structured set of real-world code snippets—both intentionally insecure and safe—that comprehensively tests each rule.

Putting It All Together

A typical workflow might look like this:

1. AST-Grep Layer: Parse code, identify known patterns, and output JSON details.
2. LLM Deep Scan: Evaluate each snippet in context, propose fixes, and rate severity.
3. Self-Verification: Ask the LLM (or multiple LLMs) to reflect on possible oversights.
4. Final Report: Consolidate AST-Grep findings, LLM suggestions, and reflect results into one comprehensive summary.

Sample Output

Once you run the toolchain, you might see an aggregated output like this (simplified for illustration):

Copy{
  "scan_results": [
    {
      "file": "insecure_logging_example.py",
      "line": 4,
      "rule_id": "no-plain-exception-logging",
      "llm_analysis": {
        "explanation": "Printing raw exception may leak sensitive data.",
        "suggested_fix": "Use a logging library with caution or hide details."
      },
      "confidence": "high"
    }
  ]
}

Implementing the DeepSeek Client

At the heart of our test fixture generation is a robust Go client for DeepSeek's AI API. This client handles all communication with DeepSeek's models, which we use to generate both vulnerable and safe code examples. Let's dive into the implementation:

Core Client Structure

First, we define our client with careful attention to configuration and error handling:

Copy// Package deepseek provides a Go client for the Deepseek API.
package deepseek

import (
    "context"
    "encoding/json"
    "fmt"
    "net/http"
    "time"
)

// Constants for API configuration
const (
    defaultBaseURL     = "https://api.deepseek.com"
    defaultModel       = "deepseek-chat-v3"
    defaultTimeout     = 30 * time.Second
    defaultTemperature = 0.0
)

// Client represents a Deepseek API client
type Client struct {
    apiKey     string
    baseURL    string
    model      string
    httpClient *http.Client
}

// NewClient creates a new Deepseek API client
func NewClient(apiKey string, opts ...ClientOption) *Client {
    if apiKey == "" {
        panic("deepseek: API key is required")
    }

    client := &Client{
        apiKey:  apiKey,
        baseURL: defaultBaseURL,
        model:   defaultModel,
        httpClient: &http.Client{
            Timeout: defaultTimeout,
        },
    }

    for _, opt := range opts {
        opt(client)
    }

    return client
}

Message Types and Request Structure

We define clear types for chat messages and API requests:

Copytype Role string

const (
    RoleSystem    Role = "system"
    RoleUser      Role = "user"
    RoleAssistant Role = "assistant"
)

type Message struct {
    Role    Role   `json:"role"`
    Content string `json:"content"`
}

type ChatCompletionRequest struct {
    Model          string    `json:"model"`
    Messages       []Message `json:"messages"`
    Temperature    float64   `json:"temperature"`
    Stream         bool      `json:"stream"`
    ResponseFormat *struct {
        Type string `json:"type"`
    } `json:"response_format,omitempty"`
}

Robust Error Handling

We implement sophisticated error handling with custom error types:

Copytype ErrorKind int

const (
    UnknownError ErrorKind = iota
    ValidationError
    AuthenticationError
    InsufficientBalanceError
    RateLimitError
    ServerError
)

type Error struct {
    kind    ErrorKind
    message string
    err     error
}

func (e *Error) Error() string {
    return fmt.Sprintf("deepseek: %s: %v", e.message, e.err)
}

Integration with the Generator

In our fixture generator, we use the DeepSeek client to create test cases:

Copyfunc (g *Generator) GenerateFixtures(rule Rule) error {
    // Create prompt for vulnerable code
    vulnRequest := &deepseek.ChatCompletionRequest{
        Messages: []deepseek.Message{
            {
                Role: deepseek.RoleUser,
                Content: fmt.Sprintf(
                    "Generate a code example that violates this rule: %s",
                    rule.Description,
                ),
            },
        },
        Temperature: 0.7, // Some randomness for variety
    }

    vulnResp, err := g.client.CreateChatCompletion(context.Background(), vulnRequest)
    if err != nil {
        return fmt.Errorf("failed to generate vulnerable fixture: %w", err)
    }

    // Parse and save the generated code
    if err := g.saveFixture(rule.ID, vulnResp.Choices[0].Message.Content, true); err != nil {
        return fmt.Errorf("failed to save vulnerable fixture: %w", err)
    }

    // Similar process for safe code example...
    return nil
}

Using the Client

To use the DeepSeek client in your own projects:

Copyfunc main() {
    // Initialize the client
    client := deepseek.NewClient("your-api-key",
        deepseek.WithModel("deepseek-chat-v3"),
        deepseek.WithHTTPClient(&http.Client{
            Timeout: 30 * time.Second,
        }),
    )

    // Create a chat completion request
    req := &deepseek.ChatCompletionRequest{
        Messages: []deepseek.Message{
            {
                Role:    deepseek.RoleUser,
                Content: "Generate a Python function with a SQL injection vulnerability",
            },
        },
        Temperature: 0.7,
    }

    // Send the request
    resp, err := client.CreateChatCompletion(context.Background(), req)
    if err != nil {
        log.Fatalf("Failed to generate code: %v", err)
    }

    // Use the response
    generatedCode := resp.Choices[0].Message.Content
    fmt.Println(generatedCode)
}

This implementation provides several key benefits:

1. Configurability: The client can be customized with different models, timeouts, and base URLs.
2. Error Handling: Comprehensive error types make debugging and error recovery straightforward.
3. Type Safety: Strong typing ensures API requests and responses are properly structured.
4. Context Support: Full context.Context support for timeout and cancellation handling.
5. Clean API: Simple, intuitive interface for making requests to DeepSeek's API.

The DeepSeek client is crucial for our fixture generation pipeline, allowing us to automatically create diverse test cases that cover both vulnerable and safe coding patterns. By using DeepSeek's AI capabilities, we can generate realistic code samples that help validate our AST-Grep rules and LLM analysis pipeline.

Takeaways

This multi-layered method greatly reduces misses and false alarms, giving security teams a clearer, more actionable report. Our approach isn't limited to one language or one LLM service, so you can adapt it across TypeScript, Python, Rust, and more—whatever your startup or enterprise environment needs. Plus, the Go-based fixture generator makes it easier to scale test-case creation across multiple rule sets and languages.

Final Thoughts

By combining the speed of AST-based scanning with the depth of AI reasoning, plus an extra reflection step, we've built a system that offers richer insights into potential vulnerabilities. Although we're just getting started, this architecture positions us for robust, scalable code security checks that can adapt to new threats over time.

Interested in leveraging this approach in your own projects? Check out TrustSight for a modern code scanning solution that integrates these techniques—and explore our DeepSeek reference implementation to see how we generate test cases and automate the scanning pipeline. Let’s protect software together.