Skip to content

SPEC.md

Latest commit

 

History

History
442 lines (363 loc) · 20.9 KB

SPEC.md

File metadata and controls

442 lines (363 loc) · 20.9 KB
Raw

AI-Powered SAST Engine for Vulnerability Detection

Project Overview

  • Project Name: VulnAI - AI-Powered Static Application Security Testing (SAST) Engine
  • Project Type: Machine Learning / Cybersecurity Product
  • Core Functionality: Train ML model to classify and detect vulnerabilities in source code, store vulnerability intelligence, and identify issues in new code snippets
  • Target Users: Security engineers, developers, DevSecOps teams

Architecture Diagram

┌─────────────────────────────────────────────────────────────────────────────┐
│                        AI-Powered SAST Engine Architecture                  │
└─────────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────────┐
│                            USER INTERFACE LAYER                             │
│  ┌─────────────────┐                    ┌────────────────────────────────┐│
│  │   REST API      │                    │      CLI Detection Tool        ││
│  │   (FastAPI)     │                    │   (Python-based Scanner)       ││
│  └────────┬────────┘                    └───────────────┬────────────────┘│
└───────────┼─────────────────────────────────────────────┼──────────────────┘
            │                                             │
            ▼                                             ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                          DETECTION ENGINE LAYER                             │
│  ┌─────────────────┐  ┌─────────────────┐  ┌────────────────────────────────┐│
│  │  Code Parser    │  │  AST Generator  │  │   Rule-Based Filter            ││
│  │  (Multi-lang)   │  │  (Tree-sitter)  │  │   (False Positive Reduction)   ││
│  └────────┬────────┘  └────────┬────────┘  └───────────────┬────────────────┘│
│           │                    │                           │                │
│           └──────────┬─────────┘                           │                │
│                      ▼                                     │                │
│  ┌──────────────────────────────────────────────────────────▼─────────────┐│
│  │                      MODEL INFERENCE ENGINE                           ││
│  │  ┌─────────────────┐  ┌─────────────────┐  ┌─────────────────────────┐││
│  │  │ CodeBERT        │  │  Similarity     │  │   Output Formatter      │││
│  │  │ Embedding       │  │  Search         │  │   (JSON Results)        │││
│  │  └────────┬────────┘  └────────┬────────┘  └──────────────┬──────────┘││
│  └───────────┼────────────────────┼─────────────────────────┼────────────┘│
└──────────────┼────────────────────┼─────────────────────────┼─────────────┘
               │                    │                         │
               ▼                    ▼                         ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                        VULNERABILITY INTELLIGENCE LAYER                     │
│  ┌─────────────────────────────────────────────────────────────────────────┐│
│  │                    PostgreSQL + pgvector Database                      ││
│  │  ┌──────────────────────┐  ┌─────────────────────────────────────────┐││
│  │  │  vulnerabilities      │  │  detected_issues                          │││
│  │  │  - id                │  │  - id                                    │││
│  │  │  - cwe_id            │  │  - file_name                             │││
│  │  │  - name              │  │  - line_number                           │││
│  │  │  - description       │  │  - detected_cwe                          │││
│  │  │  - severity          │  │  - confidence                            │││
│  │  │  - remediation       │  │  - timestamp                             │││
│  │  │  - embedding_vector  │  │  - fixed_code                            │││
│  │  └──────────────────────┘  └─────────────────────────────────────────┘││
│  └─────────────────────────────────────────────────────────────────────────┘│
└─────────────────────────────────────────────────────────────────────────────┘
               ▲
               │
┌──────────────┼────────────────────────────────────────────────────────────────┐
│              │              MODEL TRAINING PIPELINE                         │
│              ▼                                                                │
│  ┌──────────────────────────────────────────────────────────────────────────┐ │
│  │                         DATA COLLECTION LAYER                            │ │
│  │  ┌────────────┐  ┌────────────┐  ┌────────────┐  ┌────────────────────┐ │ │
│  │  │OWASP Top10 │  │ MITRE CWE  │  │   NVD CVE  │  │  CodeQL/SAST Data  │ │ │
│  │  └─────┬──────┘  └─────┬──────┘  └─────┬──────┘  └─────────┬──────────┘ │ │
│  │        └──────────────┴───────────────┴──────────────────┘              │ │
│  │                               ▼                                          │ │
│  │  ┌──────────────────────────────────────────────────────────────────────┐│ │
│  │  │                    PREPROCESSING PIPELINE                           ││ │
│  │  │  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐  ┌────────────┐ ││ │
│  │  │  │Code Cleaner │  │  Tokenizer  │  │  AST Parser│  │ Embeddings │ ││ │
│  │  │  │(Comments,   │  │  (BERT Tok) │  │ (Tree-sitter│  │(CodeBERT)  │ ││ │
│  │  │  │ Normalize)  │  │             │  │  /AST)      │  │            │ ││ │
│  │  │  └─────────────┘  └─────────────┘  └─────────────┘  └────────────┘ ││ │
│  │  └──────────────────────────────────────────────────────────────────────┘│ │
│  │                               ▼                                          │ │
│  │  ┌──────────────────────────────────────────────────────────────────────┐│ │
│  │  │                      MODEL TRAINING                                 ││ │
│  │  │  ┌─────────────────────────────────────────────────────────────────┐ ││ │
│  │  │  │              Transformer Classifier (CodeBERT)                 │ ││ │
│  │  │  │  Input: Code Snippet → Tokenize → CodeBERT → Classifier → CWE │ ││ │
│  │  │  └─────────────────────────────────────────────────────────────────┘ ││ │
│  │  └──────────────────────────────────────────────────────────────────────┘│ │
│  │                               ▼                                          │ │
│  │  ┌──────────────────────────────────────────────────────────────────────┐│ │
│  │  │                    EVALUATION & EXPORT                              ││ │
│  │  │  Precision/Recall/F1 | Confusion Matrix | Export .pt/.onnx        ││ │
│  │  └──────────────────────────────────────────────────────────────────────┘│ │
│  └──────────────────────────────────────────────────────────────────────────┘ │
└──────────────────────────────────────────────────────────────────────────────┘

Technology Stack

Core Technologies

  • Programming Language: Python 3.10+
  • ML Framework: PyTorch 2.0+
  • Transformer Model: CodeBERT, GraphCodeBERT
  • Web Framework: FastAPI
  • Database: PostgreSQL 15+ with pgvector
  • CLI: Click framework

Supporting Libraries

  • Code Parsing: Tree-sitter, python-ast
  • Data Processing: pandas, numpy
  • Model Training: transformers (HuggingFace)
  • Vector Search: pgvector, FAISS
  • API Documentation: OpenAPI/Swagger
  • Testing: pytest

Functionality Specification

1. Data Collection Module

Features:

  • Fetch vulnerability data from NVD API
  • Parse CWE database for vulnerability types
  • Collect OWASP Top 10 samples
  • Aggregate vulnerable and fixed code pairs

Data Sources:

  • NIST National Vulnerability Database (NVD)
  • MITRE Common Weakness Enumeration (CWE)
  • OWASP Foundation
  • GitHub Security Advisory Database

2. Preprocessing Pipeline

Code Cleaning:

  • Remove comments (single-line, multi-line)
  • Normalize identifiers (replace variable names)
  • Handle whitespace and formatting
  • Extract code tokens

AST-Based Parsing:

  • Generate Abstract Syntax Trees
  • Extract code structure features
  • Identify function calls, variables, data flows

Embedding Generation:

  • Use CodeBERT tokenizer
  • Generate contextual embeddings
  • Store vectors for similarity search

3. Model Architecture

Primary: CodeBERT-Based Classifier

Input: Source Code Snippet
  ↓
CodeBERT Tokenizer
  ↓
CodeBERT Encoder (12 layers)
  ↓
[CLS] Token Representation
  ↓
Dense Layer (768 → 256) + ReLU + Dropout
  ↓
Output Layer (256 → num_cwe_classes)
  ↓
Output: {is_vulnerable, cwe_id, confidence_score}

Supported Vulnerability Classes

  • CWE-89: SQL Injection
  • CWE-79: Cross-Site Scripting (XSS)
  • CWE-94: Code Injection
  • CWE-78: OS Command Injection
  • CWE-287: Insecure Authentication
  • CWE-862: Insecure Authorization
  • CWE-434: Unrestricted File Upload
  • CWE-502: Deserialization of Untrusted Data
  • CWE-119: Buffer Overflow
  • CWE-200: Information Exposure

4. Vulnerability Intelligence Storage

Database Schema

-- Vulnerabilities table
CREATE TABLE vulnerabilities (
    id SERIAL PRIMARY KEY,
    cwe_id VARCHAR(10) NOT NULL,
    name VARCHAR(255) NOT NULL,
    description TEXT,
    severity VARCHAR(20),  -- CRITICAL, HIGH, MEDIUM, LOW
    remediation TEXT,
    vulnerable_code TEXT,
    fixed_code TEXT,
    embedding_vector VECTOR(768),
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

-- Detected issues table
CREATE TABLE detected_issues (
    id SERIAL PRIMARY KEY,
    file_name VARCHAR(500),
    line_number INTEGER,
    detected_cwe VARCHAR(10),
    confidence FLOAT,
    vulnerable_code TEXT,
    suggested_fix TEXT,
    scan_timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

-- Index for vector similarity search
CREATE INDEX ON vulnerabilities USING ivfflat (embedding_vector vector_cosine_ops);

5. Detection Engine

Input: Source code snippet or file path

Processing Flow:

  1. Parse input code (auto-detect language)
  2. Generate AST for structural analysis
  3. Generate CodeBERT embeddings
  4. Run model inference
  5. Perform similarity search in vector DB
  6. Apply rule-based filters
  7. Generate results with explanations

Output Format:

{
  "is_vulnerable": true,
  "vulnerability_type": "CWE-89",
  "confidence_score": 0.92,
  "line_number": 42,
  "description": "SQL injection vulnerability detected",
  "remediation": "Use parameterized queries",
  "cwe_reference": "https://cwe.mitre.org/data/definitions/89.html"
}

6. REST API Endpoints

Method Endpoint Description
POST /api/v1/detect Detect vulnerabilities in code
GET /api/v1/vulnerabilities List stored vulnerabilities
GET /api/v1/vulnerabilities/{cwe_id} Get specific vulnerability
POST /api/v1/feedback Submit feedback for incremental learning
GET /api/v1/stats Get detection statistics
GET /api/v1/health Health check

7. CLI Tool

# Scan a single file
vulnai scan --file path/to/code.py

# Scan directory
vulnai scan --directory ./src

# Scan with specific language
vulnai scan --file app.js --language javascript

# Output JSON
vulnai scan --file main.py --output json

# Verbose mode
vulnai scan -f test.java -v

Evaluation Metrics

Primary Metrics

  • Precision: Proportion of true positives among all positive predictions
  • Recall: Proportion of actual vulnerabilities detected
  • F1-Score: Harmonic mean of precision and recall
  • False Positive Rate: Proportion of false alarms

Per-CWE Accuracy

  • Track accuracy for each vulnerability class
  • Generate confusion matrix
  • Identify model weaknesses

File Structure

ML_model_Vulnerability_Detection/
├── SPEC.md
├── TODO.md
├── requirements.txt
├── setup.py
├── vulnai/
│   ├── __init__.py
│   ├── cli/
│   │   ├── __init__.py
│   │   └── main.py
│   ├── api/
│   │   ├── __init__.py
│   │   ├── main.py
│   │   ├── routes/
│   │   │   ├── __init__.py
│   │   │   ├── detect.py
│   │   │   ├── vulnerabilities.py
│   │   │   └── feedback.py
│   │   └── models/
│   │       ├── __init__.py
│   │       └── schemas.py
│   ├── core/
│   │   ├── __init__.py
│   │   ├── config.py
│   │   └── logger.py
│   ├── detection/
│   │   ├── __init__.py
│   │   ├── engine.py
│   │   ├── parser.py
│   │   └── filter.py
│   ├── models/
│   │   ├── __init__.py
│   │   ├── classifier.py
│   │   ├── trainer.py
│   │   └── evaluator.py
│   ├── preprocessing/
│   │   ├── __init__.py
│   │   ├── cleaner.py
│   │   ├── tokenizer.py
│   │   └── embedder.py
│   ├── storage/
│   │   ├── __init__.py
│   │   ├── database.py
│   │   └── vector_store.py
│   └── data/
│       ├── __init__.py
│       ├── collector.py
│       └── loader.py
├── models/
│   ├── checkpoints/
│   └── trained/
│       └── vulnai_classifier.pt
├── data/
│   ├── raw/
│   ├── processed/
│   └── samples/
├── tests/
│   ├── __init__.py
│   ├── test_detection.py
│   ├── test_preprocessing.py
│   └── test_api.py
├── notebooks/
│   ├── data_exploration.ipynb
│   ├── model_training.ipynb
│   └── evaluation.ipynb
├── docs/
│   ├── architecture.md
│   ├── api_docs.md
│   └── usage.md
└── README.md

Implementation Phases

Phase 1: Project Setup (Week 1)

  • Initialize project structure
  • Set up development environment
  • Configure dependencies

Phase 2: Data Pipeline (Week 2-3)

  • Build data collection module
  • Implement preprocessing pipeline
  • Create embedding generation

Phase 3: Model Development (Week 4-6)

  • Design model architecture
  • Train classifier
  • Evaluate and optimize

Phase 4: Storage & API (Week 7-8)

  • Set up database
  • Build REST API
  • Implement CLI tool

Phase 5: Advanced Features (Week 9-10)

  • Multi-language support
  • False positive reduction
  • Explainability features

Phase 6: Testing & Documentation (Week 11-12)

  • Comprehensive testing
  • Documentation
  • Demo and evaluation report

Success Criteria

  1. Model Performance: F1-score > 0.85 on test set
  2. False Positive Rate: < 10% on validation set
  3. Supported Languages: Python, Java, JavaScript
  4. Detection Coverage: Top 10 OWASP vulnerabilities
  5. API Response Time: < 2 seconds per scan
  6. Model Export: .pt and .onnx formats

Security Considerations

  • All vulnerability data stored securely
  • API authentication for production use
  • Rate limiting on public endpoints
  • Input sanitization for code processing
  • Audit logging for all detections