Skip to content

Aditya-107-prog/SecureRAG

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

1 Commit
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

SecureRAG

A production-grade Retrieval-Augmented Generation (RAG) system over cybersecurity knowledge (CISA KEV + MITRE ATT&CK), enhanced with Corrective RAG (CRAG) — an LLM critic that validates retrieval quality and automatically rewrites queries when context is insufficient.


What This Project Is

Most RAG systems blindly retrieve the top-k documents and generate an answer — even if the retrieved docs are completely irrelevant. SecureRAG adds a self-correction loop:

User Query
    │
    ▼
Hybrid Search (BM25 + Vector)         ← catches both exact & semantic matches
    │
    ▼
Cross-Encoder Re-Ranker               ← reorders by true relevance score
    │
    ▼
CRAG Critic (Llama 3.3 70B)          ← judges: relevant / partial / irrelevant
    │
    ├── sufficient? ─────────────────────────────────────────────────────┐
    │                                                                     │
    └── insufficient? → Query Rewriter → Re-retrieve (max 2 attempts)    │
                                                                          ▼
                                                          Answer Generation (Llama 3.3 70B)
                                                                          │
                                                                          ▼
                                                          Answer + Cited CVE/MITRE Sources

Datasets

SecureRAG is built on two complementary public cybersecurity datasets that together cover both what vulnerabilities exist and how attackers exploit them.

1. CISA Known Exploited Vulnerabilities (KEV)

Source: https://www.cisa.gov/known-exploited-vulnerabilities-catalog Size: 1,619 CVE records Format: JSON feed, updated daily by CISA

The CISA KEV catalog is a curated list of vulnerabilities that have been actively exploited in real-world attacks — not just theoretical vulnerabilities. Every entry has been confirmed by CISA as being used by real threat actors. This makes it far higher signal than the full NVD database (200,000+ CVEs, many never exploited).

Each CVE record contains:

  • CVE ID — unique identifier (e.g., CVE-2021-44228)
  • Vendor & Product — who makes the affected software (e.g., Apache, Log4j2)
  • Vulnerability Name — human-readable name (e.g., "Apache Log4j2 Remote Code Execution Vulnerability")
  • Short Description — what the vulnerability is and how it works
  • Required Action — what organizations must do to remediate
  • Date Added — when CISA added it to the KEV catalog
  • Due Date — federal agencies' remediation deadline
  • Known Ransomware Use — whether ransomware groups have used this CVE

How we process it: Each CVE record is converted from raw JSON into a natural language prose chunk:

CVE ID: CVE-2021-44228
Vendor: Apache
Product: Log4j2
Vulnerability Name: Apache Log4j2 Remote Code Execution Vulnerability
Description: Apache Log4j2 contains a vulnerability where JNDI features do not
protect against attacker-controlled LDAP and other JNDI related endpoints...
Required Action: Apply updates per vendor instructions.
Known Ransomware Use: Known

This prose format is intentional — embedding models understand natural language better than raw JSON key-value pairs, improving retrieval quality.


2. MITRE ATT&CK Enterprise

Source: https://github.com/mitre/cti (STIX format) Size: 697 active techniques and sub-techniques Format: STIX JSON bundle (~50MB), published by MITRE

MITRE ATT&CK is a knowledge base of adversary tactics and techniques — not what vulnerabilities exist, but how attackers actually behave once they're in a system. It's organized as a matrix of Tactics (the goal) → Techniques (how to achieve it) → Sub-techniques (specific implementations).

Each technique record contains:

  • Technique ID — unique identifier (e.g., T1068, T1548.002)
  • Name — technique name (e.g., "Exploitation for Privilege Escalation")
  • Tactics — which attack phase this belongs to (e.g., Privilege Escalation, Persistence)
  • Platforms — which OS/environments are affected (Windows, Linux, macOS, Cloud)
  • Description — detailed explanation of how adversaries use this technique
  • Detection Guidance — how defenders can detect this technique in their environment

How we process it: We filter to only active (non-deprecated, non-revoked) techniques and convert each into a prose chunk:

MITRE ATT&CK Technique: T1068 - Exploitation for Privilege Escalation
Tactics: Privilege Escalation
Platforms: Containers, Linux, macOS, Windows
Description: Adversaries may exploit software vulnerabilities in an attempt
to elevate privileges...
Detection Guidance: Existing exploitation prevention tools...

Why both datasets together? CVE tells you what's broken, MITRE tells you what attackers do with it. A question like "How do adversaries escalate privileges on Windows?" is best answered by combining MITRE techniques (T1068, T1548) with real CVEs used for privilege escalation from KEV — neither dataset alone gives the complete picture.


Sample Outputs

Query: "What is CVE-2021-44228 and what product does it affect?"

Grade: RELEVANT (0 rewrites)

Answer:
CVE-2021-44228, also known as Log4Shell, is a remote code execution
vulnerability in Apache Log4j2. It allows attackers to execute arbitrary
code via JNDI features that do not protect against attacker-controlled
LDAP endpoints. The required action is to apply updates per vendor
instructions or remove affected assets from agency networks.

Sources:
  - KEV: CVE-2021-44228
  - KEV: CVE-2021-45046
  - MITRE: T1190

Query: "How do adversaries perform privilege escalation on Windows?"

Grade: RELEVANT (0 rewrites)

Answer:
Adversaries can perform privilege escalation on Windows by bypassing User
Account Control (UAC) mechanisms (T1548.002), exploiting software
vulnerabilities (T1068), or abusing elevation control mechanisms (T1548).
Real-world CVEs used for this include CVE-2019-0880 and CVE-2004-0210,
both Microsoft Windows privilege escalation vulnerabilities in the CISA KEV.

Sources:
  - MITRE: T1548.002
  - MITRE: T1068
  - MITRE: T1548
  - KEV: CVE-2019-0880
  - KEV: CVE-2004-0210

Query: "What ransomware campaigns exploited Microsoft Exchange?"

Grade: IRRELEVANT → rewrote query → PARTIAL (2 rewrites)

Rewritten query: "What specific ransomware campaigns, such as REvil or
LockBit, have leveraged CVE-2021-26855, CVE-2021-26857, CVE-2021-26858,
and CVE-2021-27065 (ProxyLogon chain) to gain initial access..."

Answer:
Ransomware campaigns have exploited the ProxyLogon exploit chain,
including CVE-2021-26855, CVE-2021-26858, CVE-2021-27065, and
CVE-2021-26857. These allow remote code execution on Microsoft Exchange
and have known ransomware use per CISA KEV.

Note: Our knowledge base contains CVE descriptions but not threat actor
attribution — for specific campaign attribution, refer to CISA advisories.

Sources:
  - KEV: CVE-2021-26855
  - KEV: CVE-2021-26858
  - KEV: CVE-2021-27065
  - KEV: CVE-2021-26857

The third example shows honest behavior — the system correctly identifies it can't fully answer the question (no threat actor data in our KB), says so, and provides what it can.


Evaluation Results

Evaluated across 10 questions comparing 3 pipeline configurations.

Config Faithfulness Answer Relevancy Context Precision
Basic RAG (vector only) 0.08 0.83 0.22
Hybrid RAG (BM25 + vector + rerank) 0.74 0.93 0.28
Hybrid + CRAG (full pipeline) 0.76 1.00 0.34

Improvement — Basic RAG → Hybrid + CRAG:

  • Faithfulness: +850% (0.08 → 0.76)
  • Answer Relevancy: +20.5% (0.83 → 1.00) — perfect score from independent judge
  • Context Precision: +54.5% (0.22 → 0.34)

Evaluation Methodology

We deliberately use three different scoring methods to avoid bias and ensure credibility. The key design principle: the model that generates answers (Groq Llama 3.3 70B) never judges its own outputs.

Faithfulness — Amazon Nova Micro (independent LLM judge)

What it measures: Is every claim in the generated answer actually supported by the retrieved documents? Low faithfulness = hallucination.

How it works: For each answer, we send the top 3 retrieved document chunks and the generated answer to Amazon Nova Micro (a completely different model from a different company — Amazon vs Meta) with this prompt:

Score the answer's faithfulness to the context:
1.0 = every claim is directly supported by the context
0.5 = most claims supported, minor unsupported details
0.0 = significant information not found in the context

Nova Micro returns a single number between 0.0 and 1.0. We average across all 10 questions.

Why Nova Micro as judge? Using the same model (Llama) to judge its own outputs creates self-serving bias — models naturally rate their own responses more favorably. Nova Micro is a different architecture with different training, making it a genuinely independent evaluator.


Answer Relevancy — Amazon Nova Micro (independent LLM judge)

What it measures: Does the answer actually address what was asked? An answer can be perfectly faithful (everything grounded in context) but still miss the question — this metric catches that failure mode.

How it works: We send the question and the generated answer to Nova Micro:

Score how well the answer addresses the question:
1.0 = directly and completely answers the question
0.5 = partially answers or addresses a related but different question
0.0 = does not answer the question at all

Why separate from faithfulness? Consider: a system retrieves a completely wrong CVE and faithfully summarizes it. Faithfulness = 1.0 (perfectly grounded in context), Relevancy = 0.0 (doesn't answer the question). Both metrics together catch all major failure modes.


Context Precision — Rule-based entity matching (zero LLM)

What it measures: What proportion of the 5 retrieved documents are actually relevant to the question?

How it works — purely with code, no LLM:

Step 1: Extract security entities from the ground truth answer and question using regex:

CVE IDs   : r'CVE-\d{4}-\d{4,7}'          # e.g. CVE-2021-44228
MITRE IDs : r'T\d{4}(?:\.\d{3})?'          # e.g. T1068, T1548.002
Keywords  : ["log4j", "exchange", "apache", "microsoft", ...]

Step 2: For each of the 5 retrieved documents, check if it contains at least one of those entities.

Step 3: Score = (documents containing expected entities) / (total documents retrieved)

Example: Question about CVE-2021-44228. Ground truth mentions CVE-2021-44228, Log4j, Apache. If 3 of 5 retrieved docs contain these → score = 3/5 = 0.60.

Why rule-based? Context precision is about whether specific identifiers appear in retrieved text — regex is more accurate and reproducible than asking an LLM to make this judgment. Zero cost, zero rate limiting, fully deterministic, same score every run.

Known limitation: Context precision scores are relatively low (0.22-0.34) across all configs because our 5-document retrieval window doesn't always surface every relevant CVE mentioned in the ground truth. Increasing top-k or expanding the knowledge base would improve this metric.


Why three different methods?

Metric Method Reason
Faithfulness Nova Micro (LLM) Requires reading comprehension — regex can't check if claims are grounded
Answer Relevancy Nova Micro (LLM) Requires semantic understanding of question intent
Context Precision Rule-based regex Checking if specific IDs appear in text — regex is more accurate and objective than LLM

This design means no single model or method dominates the evaluation — each metric is scored by the most appropriate tool for that specific measurement.


Tech Stack

Component Technology Cost
LLM (critic + generation) Groq — Llama 3.3 70B Free
LLM fallback + eval judge Amazon Nova Micro (AWS Bedrock) ~$0.00
Embeddings BAAI/bge-base-en-v1.5 (local) Free
Vector store ChromaDB (persistent, local) Free
Keyword search BM25 via rank-bm25 Free
Rank fusion Reciprocal Rank Fusion (RRF)
Re-ranker cross-encoder/ms-marco-MiniLM-L-6-v2 Free
Context Precision eval Rule-based regex entity matching Free
UI Streamlit Free
Total cost ~$0.00

Project Structure

security-rag-crag/
│
├── data/                            # all data (mostly gitignored)
│   ├── raw/                         # raw downloads from CISA/MITRE (gitignored)
│   ├── processed/                   # cleaned JSONL chunks (gitignored)
│   │   ├── cisa_kev_chunks.jsonl    # 1,619 CVE prose chunks
│   │   └── mitre_attack_chunks.jsonl # 697 technique prose chunks
│   ├── chroma_db/                   # ChromaDB vector index (gitignored, ~500MB)
│   ├── bm25_cache/                  # BM25 pickle indexes (gitignored)
│   ├── eval_questions.json          # 10 hand-crafted Q&A pairs 
│   └── eval_results.json            # evaluation scores 
│
├── src/                             # all source code
│   │
│   ├── llm_client.py                # central LLM client
│   │                                # → Groq (Llama 3.3 70B) as primary
│   │                                # → Nova Micro (Bedrock) as fallback
│   │                                # → built-in token + cost tracking
│   │
│   ├── ingest/                      # data fetching + processing
│   │   ├── fetch_cisa_kev.py        # downloads CISA KEV JSON → prose chunks
│   │   └── fetch_mitre.py           # downloads MITRE ATT&CK bundle → prose chunks
│   │
│   ├── retrieval/                   # search layer
│   │   ├── vector_store.py          # builds ChromaDB with BGE embeddings
│   │   │                            # → separate collections for KEV + MITRE
│   │   │                            # → persistent to disk, loads instantly after first run
│   │   ├── hybrid_search.py         # BM25 + vector search fused with RRF
│   │   │                            # → BM25 catches exact CVE ID matches
│   │   │                            # → vector search catches semantic matches
│   │   │                            # → RRF boosts docs that rank well in both
│   │   └── reranker.py              # cross-encoder re-ranker
│   │                                # → fetches 10 candidates, keeps top 5
│   │                                # → reads query+doc together for accurate scoring
│   │
│   ├── crag/                        # corrective RAG layer
│   │   └── critic.py                # retrieval quality grader
│   │                                # → grades: relevant / partial / irrelevant
│   │                                # → rewrites query if grade is not relevant
│   │                                # → max 2 rewrite attempts before generating anyway
│   │
│   ├── generation/                  # answer generation
│   │   └── rag_pipeline.py          # full end-to-end pipeline
│   │                                # → hybrid search → rerank → CRAG → generate
│   │                                # → returns answer + sources + grade + rewrite log
│   │
│   └── eval/                        # evaluation
│       └── ragas_eval.py            # independent evaluation suite
│                                    # → generator : Groq Llama 3.3 70B
│                                    # → judge     : Amazon Nova Micro (independent)
│                                    # → precision : rule-based entity matching
│
├── app/
│   └── streamlit_app.py             # interactive demo UI
│                                    # → mode toggle: Basic RAG vs CRAG
│                                    # → shows retrieved sources + CRAG grade
│                                    # → query rewrite log panel
│                                    # → session history
│
├── .env                             # API keys (never committed)
├── .env.example                     # template for new users
├── .gitignore                       # excludes data/, chroma_db/, .env
├── requirements.txt
└── README.md

How to Run

# 1. Clone and install
git clone https://github.com/yourusername/security-rag-crag
cd security-rag-crag
pip install -r requirements.txt

# 2. Set up environment
cp .env.example .env
# Edit .env and add your GROQ_API_KEY

# 3. Fetch datasets (run once)
python -m src.ingest.fetch_cisa_kev    # ~2 seconds, 1619 CVEs
python -m src.ingest.fetch_mitre       # ~60 seconds, 50MB download

# 4. Build vector store (run once, ~15 mins on CPU)
python -m src.retrieval.vector_store

# 5. Run full pipeline
python -m src.generation.rag_pipeline

# 6. Run evaluation (3 configs x 10 questions, ~10 mins)
python -m src.eval.ragas_eval

# 7. Launch Streamlit UI
streamlit run app/streamlit_app.py

Environment Variables

# .env.example — copy to .env and fill in values
GROQ_API_KEY=your_groq_api_key_here     # get free key at console.groq.com
AWS_REGION=ap-southeast-2               # needed for Nova Micro eval judge

Build Log — Issues Faced & How We Solved Them

Issue 1: AWS Bedrock — Anthropic Models Blocked

Problem: Anthropic Claude models on Bedrock threw AccessDeniedException: INVALID_PAYMENT_INSTRUMENT — the AWS free plan blocks AWS Marketplace model subscriptions required for Anthropic models.

Solution: Switched primary LLM to Groq (free, Llama 3.3 70B) and used Amazon Nova Micro as fallback — Nova Micro doesn't require Marketplace approval. Built a unified llm_client.py with automatic fallback.

Lesson: Always validate LLM provider access at the start of a project before building on top of it.


Issue 2: Groq Model Decommissioned Mid-Build

Problem: llama-3.1-70b-versatile was decommissioned mid-build and threw model_decommissioned error.

Solution: Updated to llama-3.3-70b-versatile — stronger reasoning and better structured output generation.

Lesson: Pin model versions in production. In development, check provider deprecation notices when errors appear.


Issue 3: Bedrock Inference Profile IDs

Problem: Direct model IDs threw ValidationException: model identifier is invalid — newer Bedrock models require region-specific inference profile IDs.

Solution: Listed profiles programmatically via client.list_inference_profiles() to get correct au. prefixed IDs for Sydney region.

Lesson: Always list available inference profiles programmatically rather than hardcoding model IDs from documentation.


Issue 4: RAGAS Dependency Hell

Problem: RAGAS had broken dependencies (langchain_community.chat_models.vertexai missing), conflicting langchain-core versions across 6+ packages, and secretly required OpenAI API keys despite claiming to be framework-agnostic.

Final solution: Replaced RAGAS with a custom evaluation suite — Nova Micro as independent LLM judge for faithfulness/relevancy, rule-based entity matching for context precision. More transparent, zero hidden dependencies, fully explainable.

Lesson: Popular evaluation frameworks often have hidden dependencies on paid APIs. A custom evaluation suite gives full control and is more defensible in interviews.


Issue 5: Self-Evaluation Bias

Problem: Initial eval used Groq Llama 3.3 70B for both generation and judging — the same model rating its own outputs inflates scores artificially.

Solution: Separated generator from judge: Groq Llama 3.3 70B generates answers, Amazon Nova Micro (different company, different architecture) independently judges faithfulness and relevancy, rule-based regex objectively scores context precision.

Lesson: Evaluation credibility requires independence between generator and judge. Using the same model for both is a methodological flaw that undermines all results.


Issue 6: Python Module Import Paths

Problem: ModuleNotFoundError: No module named 'retrieval' when running with python -m.

Solution: Changed all imports to absolute paths from project root (from src.retrieval.vector_store import ...) and added sys.path.insert(0, project_root) at top of each file.

Lesson: Always use absolute imports from the project root in multi-module Python projects.


Issue 7: ChromaDB First-Run Embedding Time

Problem: Embedding 2,316 documents took ~13 minutes on CPU on first run.

Solution: ChromaDB PersistentClient caches to disk — subsequent runs load in seconds. BM25 indexes also cached as pickle files.

Lesson: Always use persistent storage for embedding jobs so you pay the cost only once.


Resume Bullet Points

  • Built SecureRAG, a production-grade cybersecurity Q&A system over 2,316 CISA KEV + MITRE ATT&CK documents using hybrid BM25+vector retrieval with Reciprocal Rank Fusion and cross-encoder re-ranking
  • Implemented Corrective RAG (CRAG) with LLM-based retrieval quality grading and automatic query rewriting, improving answer faithfulness by 850% and achieving perfect answer relevancy (1.00) over a baseline RAG system
  • Designed a rigorous 3-method evaluation framework: Amazon Nova Micro as independent LLM judge (different architecture from generator), rule-based CVE/MITRE entity matching for objective context precision — deliberately separating generator from evaluator to eliminate self-evaluation bias
  • Built entirely on free infrastructure: Groq (Llama 3.3 70B), local BGE embeddings, ChromaDB — total LLM cost $0.00

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors

Languages