AI-Assisted Vulnerability Prioritization and Remediation Suggestion Tool
Introduction
In today’s threat landscape, organizations face an increasing volume of vulnerabilities. Traditional vulnerability management methods are time-consuming and lack intelligent prioritization and automated remediation.
This project aims to build an AI-powered system that:
- Automates vulnerability scanning, classification, and remediation recommendation.
- Uses tools like Nmap and OpenVAS for scanning.
- Applies NLP and Machine Learning to classify and suggest fixes for vulnerabilities.
Project Objective
To develop an AI-assisted vulnerability management tool that:
- Automates identification, classification, and prioritization of vulnerabilities.
- Assigns severity levels.
- Recommends remediation steps using real-world data and AI.
Tools & Technologies
- Python
- Nmap, OpenVAS
- scikit-learn
- spaCy, NLTK, SBERT, T5, DistilBART
- matplotlib, seaborn, Streamlit
- VS Code
Scope
Focus on detecting and analyzing network and OS vulnerabilities, integrating data from:
- NVD CVE feeds
- OpenVAS & Nmap scans
- Exploit-DB repositories
Uses AI to classify vulnerabilities and suggest remediation steps.
Methodology
Data Collection
- 5 years of CVE data from NVD JSON feeds.
- Parsed OpenVAS & Nmap scan results.
- Scraped Exploit-DB for exploit availability.
Data Preprocessing & Feature Engineering
- Cleaned descriptions using spaCy and NLTK.
- Extracted features: CVSS score, attack vector, exploit availability.
- Generated embeddings using TF-IDF and SBERT.
AI Model Training
- Trained classifiers: Logistic Regression, Random Forest, fine-tuned BERT.
- Output labels: Low, Medium, High, Critical.
Integration & Interface
- Automated pipeline for scan result ingestion and AI processing.
- GUI using Streamlit for interactive usage.
Reporting & Visualization
- Dashboard with real-time filtering.
- CVE severity breakdown, exploit availability, patch coverage.
Results & Evaluation
Vulnerability Severity Classification:
| Model | Accuracy | Precision | Recall | F1-Score | |:—————– |:———-|:———-|:——–|:———-| | Logistic Regression | 0.76 | 0.74 | 0.73 | 0.73 | | Random Forest | 0.81 | 0.80 | 0.79 | 0.79 | | BERT (fine-tuned) | 0.88 | 0.86 | 0.85 | 0.85 |
Scan Result Aggregation:
- CVE match rate: 91%
- False positives: 7%
- False negatives: 4%
Remediation Suggestion Module:
- Direct URL mapping from remediation datasets.
- TF-IDF + NearestNeighbors for semantic suggestions.
- Gemini LLM used for generating detailed remediations.
Challenges Faced
- OpenVAS Setup on Limited Hardware
- Class Imbalance in CVE Severity Data
- Ambiguous or Missing Remediation Steps
- Slow BERT Fine-Tuning due to Hardware Constraints
Key Achievements
- Integrated toolchain with OpenVAS, Nmap, CVE datasets.
- Built AI models with BERT achieving 88% accuracy.
- Developed a semantic remediation engine.
- Created interactive dashboards using Streamlit.
Practical Applications
- Enterprise vulnerability management.
- Automated DevSecOps pipelines.
- Security audit reporting and prioritization.
Source Code
- You can view and explore the full source code for this project on GitHub.
Future Enhancements
- Dockerization for easy deployment.
- REST API integration.
- Live threat feed integration.
- Transformer-based remediation generation.
- Custom alerts for high-severity vulnerabilities.
References
Project by: [Fahad Raza] (i221768@nu.edu.pk), Me , [Syed Inshal Yousaf] (i227439@nu.edu.pk)
Date: 05-04-2025