An innovative AI-powered learning assistant for solving cybersecurity challenges, leveraging advanced AI models like White Rabbit with Retrieval-Augmented Generation (RAG) and Hypothetical Document Embeddings (HyDE) for accurate, context-aware responses.
- About the Project
- System Architecture
- Prerequisites
- Setup Instructions
- Usage
- Directory Structure
- Acknowledgments
The Generative Cybersecurity Application is designed to enhance learning and problem-solving in cybersecurity using generative AI models. It offers:
- Context-aware assistance for cybersecurity challenges.
- Ethical educational support using authorized materials like OverTheWire Bandit wargame series.
- Integration of AI-driven query enhancement and retrieval-based learning.
- Query Enhancement: Improves user queries using LLM-based processing.
- Generative AI Responses: Provides accurate answers using White Rabbit.
- Vector Database Integration: Retrieves relevant cybersecurity documentation.
- User-Friendly Frontend: A Streamlit UI for an intuitive experience.
- Educational Focus: Covers fundamental Linux/Bash commands, SSH connectivity, file analysis, and basic cryptography.
The system is built on a modular architecture, comprising:
- Generative Module API:
- Backend for AI model inference and FastAPI endpoints.
- Streamlit Module API:
- Frontend providing an interactive user interface.
System Diagram: Copy and paste the code from SystemDiagram.md
into Mermaid Live Editor to visualize the architecture.
Ensure the following are installed on your system:
- Docker Engine (19.03.0+)
- Docker Compose (2.0+)
- NVIDIA GPU with CUDA support
- NVIDIA Container Toolkit (nvidia-docker2)
Installing Docker
Run the following command to remove any previously installed Docker-related packages that might conflict:
for pkg in docker.io docker-doc docker-compose docker-compose-v2 podman-docker containerd runc; do sudo apt-get remove $pkg; done
-
Update the
apt
package index:sudo apt-get update
-
Install required packages:
sudo apt-get install ca-certificates curl
-
Create the
/etc/apt/keyrings
directory:sudo install -m 0755 -d /etc/apt/keyrings
-
Download and add Dockerβs GPG key:
sudo curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc sudo chmod a+r /etc/apt/keyrings/docker.asc
-
Add the Docker repository:
echo "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
-
Update the package index:
sudo apt-get update
Install Docker Engine, CLI, and associated plugins:
sudo apt-get install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin
Run the following command to verify that Docker is installed correctly:
sudo docker run hello-world
-
Add your user to the Docker group:
sudo groupadd docker sudo usermod -aG docker $USER
-
Apply the changes:
newgrp docker
-
Verify that you can run Docker commands without
sudo
:docker run hello-world
To start Docker and containerd on system boot:
sudo systemctl enable docker.service
sudo systemctl enable containerd.service
Installing NVIDIA Container Toolkit
Run the following command to add NVIDIA's GPG key and repository:
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
Edit the repository file to enable experimental components:
sudo sed -i -e '/experimental/ s/^#//g' /etc/apt/sources.list.d/nvidia-container-toolkit.list
sudo apt-get update
sudo apt-get install -y nvidia-container-toolkit
Set NVIDIA as the default runtime for Docker:
sudo nvidia-ctk runtime configure --runtime=docker
sudo systemctl restart docker
-
Check that the NVIDIA Container Toolkit is installed:
dpkg -l | grep nvidia-container-toolkit
Expected Output:
ii nvidia-container-toolkit 1.16.2-1 amd64 NVIDIA Container toolkit ii nvidia-container-toolkit-base 1.16.2-1 amd64 NVIDIA Container Toolkit Base
-
Verify NVIDIA runtime is available to Docker:
sudo docker info | grep Runtimes
Expected Output:
Runtimes: io.containerd.runc.v2 nvidia runc
git clone https://github.com/iamamiramine/generative-cybersecurity.git
cd generative-cybersecurity
Refer to this section in the README for the installation commands.
docker network create generative-cybersecurity-network
Ensure the following folders exist:
models/
- Stores AI models.data/
- Stores application data (only.txt
files are supported).
Use the following command to download the required model:
curl -X POST http://localhost:7575/download_llm -H "Content-Type: application/json" -d '{"model_name": "WhiteRabbitNeo/WhiteRabbitNeo-2.5-Qwen-2.5-Coder-7B"}'
Alternatively, use Swagger UI at http://localhost:7575/docs#/
.
docker compose build
docker compose up
- Streamlit Frontend: http://localhost:7500
- Generative API Swagger UI: http://localhost:7575/docs
- Without Context:
POST /load_model
POST /load_pipeline
POST /load_chain
POST /generate
- With Context:
POST /load_model
POST /load_pipeline
POST /load_docs
POST /load_ensemble_retriever_from_docs
POST /load_chain
POST /generate
generative-cybersecurity/
βββ models/ # Contains AI models
βββ data/ # Contains application data
βββ shared/ # Shared resources
β βββ assets/ # Images and other assets
β βββ config/ # Configuration files
βββ generative-module-api/ # Backend service
βββ streamlit-module-api/ # Frontend service
- Load the model.
- Load the pipeline.
- Load the documents.
- Load the ensemble retriever.
- Load the chain.
- Generate text.
The UI dynamically reloads components like the pipeline and chain when parameters or context change.
Contributions are what make the open source community such an amazing place to learn, inspire, and create. Any contributions you make are greatly appreciated.
If you have a suggestion that would make this better, please fork the repo and create a pull request. You can also simply open an issue with the tag "enhancement". Don't forget to give the project a star! Thanks again!
- Fork the Project
- Create your Feature Branch (
git checkout -b feature/AmazingFeature
) - Commit your Changes (
git commit -m 'Add some AmazingFeature'
) - Push to the Branch (
git push origin feature/AmazingFeature
) - Open a Pull Request
Special thanks to:
- The White Rabbit open-source community for their LLM contributions.
- The OverTheWire community for providing ethical cybersecurity challenges.
- Open-source tools and libraries that made this project possible.