How AI-Infra-Guard Performs Fingerprinting for AI Frameworks: A Code-Level Breakdown

AI-Infra-Guard performs AI framework fingerprinting by loading YAML template files from data/fingerprints, parsing them into a custom state machine, and matching them against HTTP responses at runtime, supplemented by hard-coded preloaded detectors for well-known services.

AI-Infra-Guard is an open-source security scanner from Tencent designed to identify AI infrastructure and services. Its fingerprinting engine discovers which AI frameworks—such as LLaMA-CPP, MLflow, and Ollama—are running on a target by combining declarative YAML rules with custom Go parsers and preload detectors. Understanding how AI-Infra-Guard performs fingerprinting for AI frameworks reveals a hybrid approach that balances flexible template-driven scanning with high-performance, hard-coded recognition.

Fingerprint Template Loading and Initialization

The fingerprinting workflow begins at startup when the CLI or web server initializes the engine. The --fps flag specifies the template directory, defaulting to data/fingerprints, which is scanned recursively for every *.yaml file.

Each YAML file is deserialized into a FingerPrint struct defined in common/fingerprints/parser/parser.go. According to the Tencent/AI-Infra-Guard source code, this struct holds the rule definitions that drive the entire matching process.

YAML Parsing and State Machine Evaluation

Once loaded, the parser in common/fingerprints/parser builds a state machine from each template. The implementation supports token streams, syntax trees, and a stack-based evaluator to process large HTTP response bodies efficiently.

These components are implemented across the parser package, with core token logic residing in common/fingerprints/parser/token.go. This architecture allows complex response regexes, header checks, and request patterns to be evaluated without excessive memory overhead.

Pre-loaded Detectors for High-Speed Recognition

In addition to YAML templates, the engine leverages hard-coded detectors for popular services. These reside in the preload package—such as common/fingerprints/preload/mlflow.go and the corresponding llama-cpp.go detector—and implement the FingerprintDetector interface.

Pre-loaded detectors provide fast, deterministic checks for common endpoints that are known to expose framework-specific artifacts. At runtime, they are appended to the template slice via preload.CollectedFpReqs() before scanning begins.

Runtime Matching in the Scan Runner

The runner in common/runner/runner.go orchestrates the actual scan. It collects all fingerprint objects, both from parsed YAML templates and pre-loaded detectors, and iterates over the target URLs.

For each target, the runner issues the HTTP requests defined in the fingerprint rules, feeds the resulting responses back into the parser’s evaluator, and records any successful matches. Detected fingerprints are stored in the scan result under the fingerprints field, as defined in common/runner/result.go.

WebSocket API for Fingerprint Management

AI-Infra-Guard exposes a WebSocket knowledge API that supports CRUD operations on fingerprint definitions at /api/v1/knowledge/fingerprints. Handlers in common/websocket/knowledge_api.go load, edit, and delete YAML files directly in the data/fingerprints directory, enabling dynamic updates without requiring a full rebuild.

Code Example: Initializing and Running Fingerprint Checks

The following Go snippet demonstrates how the engine bootstraps its fingerprint list and applies it during a scan task:

// Initialize the fingerprint engine (run by CLI/web server)
func initFingerprints(dir string) ([]parser.FingerPrint, error) {
    fps, err := parser.LoadFromDir(dir)   // parses all *.yaml* under data/fingerprints
    if err != nil {
        return nil, err
    }
    // Add pre-loaded detectors (e.g., MLflow, LLaMA-CPP)
    fps = append(fps, preload.CollectedFpReqs()...)
    return fps, nil
}

// Example usage inside a scan task
func scanTarget(url string, fps []parser.FingerPrint) []preload.FpResult {
    var results []preload.FpResult
    for _, fp := range fps {
        if fp.Match(url) {               // sends HTTP request(s) defined in the fingerprint
            results = append(results, preload.FpResult{
                Name: fp.Info.Name,
                Url:  url,
            })
        }
    }
    return results
}

In this flow, parser.LoadFromDir() handles the YAML ingestion from paths such as data/fingerprints/llama-cpp.yaml, while preload.CollectedFpReqs() injects the hard-coded detectors for services like MLflow.

Summary

• AI-Infra-Guard uses a hybrid fingerprinting engine that combines recursive YAML template loading with hard-coded preload detectors.
• The common/fingerprints/parser package transforms templates into an efficient state machine backed by token streams and a stack-based evaluator.
• Runtime scanning is driven by common/runner/runner.go, which executes HTTP requests and populates the fingerprints field in the result object.
• An administrative WebSocket API at /api/v1/knowledge/fingerprints allows operators to manage rules dynamically.

Frequently Asked Questions

What file format does AI-Infra-Guard use for fingerprint templates?

AI-Infra-Guard stores fingerprint templates as YAML files inside the data/fingerprints directory. Each file defines request patterns, response regexes, and header checks that the parser converts into an evaluable rule set.

How does AI-Infra-Guard match fingerprints against a target URL?

During a scan, the runner in common/runner/runner.go iterates over all loaded fingerprints. For each target URL, it issues the configured HTTP requests and passes the responses to the parser. The parser’s state machine, built from tokens and syntax trees in common/fingerprints/parser/token.go, evaluates whether the response satisfies the template conditions.

What is the difference between YAML fingerprints and pre-loaded detectors?

YAML fingerprints are declarative rules loaded from disk that offer flexibility for new or custom frameworks. Pre-loaded detectors are Go implementations of the FingerprintDetector interface—such as those in common/fingerprints/preload/mlflow.go—that provide optimized, hard-coded checks for well-known AI services.

Can fingerprint definitions be modified without restarting AI-Infra-Guard?

Yes. The WebSocket knowledge API exposes the /api/v1/knowledge/fingerprints endpoint, which handlers in common/websocket/knowledge_api.go use to create, update, and delete YAML files directly in the data/fingerprints directory. This design enables runtime updates to the rule base without restarting the service.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}