This repository contains the code accompanying the paper: Detecting malicious PDF using CNN. It is implemented using PyTorch.
To set up just create a virtual environment with python3 and run:
pip install -r requirements.txt
In order to train a model, you need to have:
- The pdfs files you want to train on in a local folder
- A csv containing about information about the files. The format should be the same as samples.csv
In order to reproduce our experiments you can:
- Download the list of files given in training_files.csv from VirusTotal
- Download the contagio dump for PDF files and use training_contagio.csv
To run a training, run the file train.py
python3 train.py ModelB training_files.csv data/pdfs/ --name training1 --gpu cuda:3
- The model in the folder trainings/. It should be loaded with torch.load().
- The logs in the folder logs/.
- A PNG file containing the ROC on the test set in the current working directory.
usage: train.py [-h] [--name NAME] [--gpu GPU] [--resample] [--cont]
[--contagio]
model files_csv data_path
positional arguments:
model Model to use, should be either 'ModelA', 'ModelB', or 'ModelC'
files_csv CSV containing the files for the training and some info. Format
should be the same as sample.csv
data_path Directory in which the files are stored, the name of the files
must be to the hash in the csv file.
optional arguments:
-h, --help show this help message and exit
--name NAME Name of the training (for the log file, the model object and
the ROC picture)
--gpu GPU Which GPU to use, default will be cuda:0
--resample Whether to resample the train set
--cont Whether to continue old training
--contagio Split train test for contagio dataset