ma_cisco_malware/arguments.py

103 lines
3.8 KiB
Python

import argparse
import os
parser = argparse.ArgumentParser()
parser.add_argument("--mode", action="store", dest="mode",
default="")
parser.add_argument("--train", action="store", dest="train_data",
default="data/full_dataset.csv.tar.gz")
parser.add_argument("--data", action="store", dest="train_data",
default="data/full_dataset.csv.tar.gz")
parser.add_argument("--test", action="store", dest="test_data",
default="data/full_future_dataset.csv.tar.gz")
parser.add_argument("--model", action="store", dest="model_path",
default="results/model_x")
parser.add_argument("--models", action="store", dest="model_paths", nargs="+",
default=[])
parser.add_argument("--type", action="store", dest="model_type",
default="final") # inter, final, staggered
parser.add_argument("--depth", action="store", dest="model_depth",
default="small") # small, medium
parser.add_argument("--model_output", action="store", dest="model_output",
default="both")
parser.add_argument("--batch", action="store", dest="batch_size",
default=64, type=int)
parser.add_argument("--epochs", action="store", dest="epochs",
default=10, type=int)
# parser.add_argument("--samples", action="store", dest="samples",
# default=100000, type=int)
#
# parser.add_argument("--samples_val", action="store", dest="samples_val",
# default=10000, type=int)
#
parser.add_argument("--embd", action="store", dest="embedding",
default=128, type=int)
parser.add_argument("--hidden_char_dims", action="store", dest="hidden_char_dims",
default=256, type=int)
parser.add_argument("--window", action="store", dest="window",
default=10, type=int)
parser.add_argument("--domain_length", action="store", dest="domain_length",
default=40, type=int)
parser.add_argument("--domain_embd", action="store", dest="domain_embedding",
default=512, type=int)
# parser.add_argument("--queue", action="store", dest="queue_size",
# default=50, type=int)
#
# parser.add_argument("--p", action="store", dest="p_train",
# default=0.5, type=float)
#
# parser.add_argument("--p_val", action="store", dest="p_val",
# default=0.01, type=float)
#
# parser.add_argument("--gpu", action="store", dest="gpu",
# default=0, type=int)
#
# parser.add_argument("--tmp", action="store_true", dest="tmp")
#
parser.add_argument("--stop_early", action="store_true", dest="stop_early")
parser.add_argument("--balanced_weights", action="store_true", dest="class_weights")
parser.add_argument("--gpu", action="store_true", dest="gpu")
parser.add_argument("--new_model", action="store_true", dest="new_model")
def get_model_args(args):
return [{
"model_path": model_path,
"embedding_model": os.path.join(model_path, "embd.h5"),
"clf_model": os.path.join(model_path, "clf.h5"),
"train_log": os.path.join(model_path, "train.log.csv"),
"train_h5data": args.train_data + ".h5",
"test_h5data": args.test_data + ".h5",
"future_prediction": os.path.join(model_path, f"{os.path.basename(args.test_data)}_pred.h5")
} for model_path in args.model_paths]
def parse():
args = parser.parse_args()
args.embedding_model = os.path.join(args.model_path, "embd.h5")
args.clf_model = os.path.join(args.model_path, "clf.h5")
args.train_log = os.path.join(args.model_path, "train.log.csv")
args.train_h5data = args.train_data + ".h5"
args.test_h5data = args.test_data + ".h5"
args.future_prediction = os.path.join(args.model_path, f"{os.path.basename(args.test_data)}_pred.h5")
return args