import os

from secml_malware.attack.blackbox.c_wrapper_phi import CEmberWrapperPhi
from secml_malware.models import CClassifierEmber

from counterfit.core import config
from counterfit.core.targets import SecMLMalwareTarget


class GBDTTarget(SecMLMalwareTarget):
	model_name = 'gbdt-ember'
	model_endpoint = os.path.join(config.targets_path, 'pe_malware', 'pretrained', 'ember_model.txt')

	def __init__(self):
		super(GBDTTarget, self).__init__()
		if not os.path.isfile(self.model_endpoint):
			raise FileNotFoundError(f"Can't load ember gbdt model, file not found at {self.model_endpoint}")
		model = CClassifierEmber(self.model_endpoint)
		self.model = CEmberWrapperPhi(model)
