Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added PDF malware Detection Pipeline #205

Merged
merged 1 commit into from
Oct 10, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10,027 changes: 10,027 additions & 0 deletions Prediction Models/PDF_Malware_Detection/Dataset/PDFMalware2022.csv

Large diffs are not rendered by default.

62 changes: 62 additions & 0 deletions Prediction Models/PDF_Malware_Detection/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
# PDF Malware Detection

This project implements a machine learning-based system for detecting potential malware in PDF files. It includes feature extraction from PDF files, model training, and a prediction script for classifying PDFs as potentially malicious or benign.

## Components

1. **Feature Extraction** (`pdf_feature_extraction.py`)
- Extracts various features from PDF files using PyMuPDF and pdfid.
- Features include metadata, structural elements, and presence of potentially risky elements.

2. **Model Training** (`pdf_malware_dataset_training.py`)
- Prepares the dataset, handles data cleaning and preprocessing.
- Trains a Random Forest classifier for malware detection.
- Includes code for hyperparameter tuning (commented out).

3. **Prediction Script** (`predict_malware.py`)
- Uses the trained model to predict whether a given PDF file is potentially malicious.

## Setup

1. Install required dependencies:
```
pip install numpy pandas matplotlib scikit-learn imblearn PyMuPDF pdfid joblib
```

2. Ensure you have the dataset file `PDFMalware2022.csv` in the `Dataset` folder.

## Usage

### Training the Model

1. Run the `pdf_malware_dataset_training.py` script to train the model:
```
python pdf_malware_dataset_training.py
```
This will create a `random_forest_model.pkl` file containing the trained model.

### Predicting Malware

1. Use the `predict_malware` function in `predict_malware.py` to classify a PDF file:
```python
from predict_malware import predict_malware

result = predict_malware("path/to/your/pdf_file.pdf")
print("Prediction (0: Benign, 1: Malicious):", result)
```

2. Alternatively, run the script directly:
```
python predict_malware.py path/to/your/pdf_file.pdf
```

## Note

This project is for educational and research purposes only. It should not be used as a sole means of determining file safety. Always use caution when dealing with potentially malicious files and consult with cybersecurity professionals for comprehensive security measures.

## Future Improvements

- Implement more advanced feature extraction techniques.
- Explore other machine learning algorithms for potentially better performance.
- Add a user-friendly interface for easier interaction with the prediction system.
- Incorporate regular model updates with new malware samples to keep the detection current.
Loading
Loading