def do_set_atk(args):
if global_state.attack is None:
error_prompt('You must first define an attack to set its parameters.')
return
atk = global_state.attack
if isinstance(atk, CEnd2EndMalwareEvasion):
if args.key in wb_atk_map:
atk_type = type(global_state.attack)
if atk_type.__name__ not in wb_atk_map[args.key]:
if not hasattr(atk, args.key):
error_prompt(f'Key {args.key} not defined for object.')
error_prompt(f'You can set: {wb_atk_map.keys()}.')
setattr(atk, args.key, args.value)
success_prompt(f'Correctly set {args.key} <- {args.value}')
else:
setattr(atk, wb_atk_map[args.key][atk_type.__name__],
args.value)
success_prompt(f'Correctly set "{args.key}" to {args.value}')
else:
error_prompt(
f'Unable to set "{args.key}". Key not found or not-editable from here.'
)
elif isinstance(atk, CBlackBoxProblem):
if args.key in bb_atk_map:
atk_type = type(global_state.attack)
setattr(atk, bb_atk_map[args.key], args.value)
success_prompt(f'Correctly set {args.key} <- {args.value}')
else:
error_prompt(
'Unable to set {args.key}. Key not found or not-editable from here.'
)
def blackbox(args):
if global_state.target is None:
error_prompt('You have first to set a target.')
return
if args.type is None:
error_prompt('You have to set an attack type.')
error_prompt(f'Chose from this list: {BYTE_ATTACKS + GAMMA_ATTACKS}')
return
if args.inject is None:
if args.type != PARTIAL_DOS:
error_prompt('You have to set an injection amount.')
return
else:
args.inject = 58
if 'gamma' in args.type:
if args.goodware_folder is None:
if global_state.goodware_folder is None:
error_prompt(
'GAMMA needs to harvest samples from goodware, set --goodware_folder.'
)
return
args.goodware_folder = global_state.goodware_folder
args.model = create_wrapper_for_global_target()
attack = create_byte_based_black_box_attack(
args) if 'gamma' not in args.type else create_gamma_black_box_attack(
args)
global_state.attack = attack
success_prompt(f'Set up attack: {args.type}')
def create_wrapper_for_global_target():
if type(global_state.target) == CClassifierEmber:
return CEmberWrapperPhi(global_state.target)
if type(global_state.target) == CClassifierEnd2EndMalware:
return CEnd2EndWrapperPhi(global_state.target)
if type(global_state.target) == CClassifierSorel:
return CEmberWrapperPhi(global_state.target)
if hasattr(global_state.target, 'load_wrapper'):
try:
return global_state.target.load_wrapper()
except Exception as e:
crash_prompt("Error in loading wrapper of plugin model!")
crash_prompt(f"Exception was {e}")
raise e
error_prompt('Incorrect target')
raise NotImplementedError('Incorrect target')
def data(args):
path = args.path
if path is None:
error_prompt('You have to set a path to a file or folder.')
return
if not os.path.isfile(path) and not os.path.isdir(path):
error_prompt('{path} does not point to a file or folder.')
return
if args.goodware:
if os.path.isdir(args.path):
global_state.goodware_folder = args.path
success_prompt('Goodware folder path correctly loaded!')
return
error_prompt(
"Goodware must be specified as a folder, not single files!")
return
if os.path.isfile(path):
file_list = [path]
else:
file_list = sorted([os.path.join(path, f) for f in os.listdir(path)])
if args.magic:
file_list = [f for f in file_list if args.magic in magic.from_file(f)]
if args.contains is not None:
file_list = [f for f in file_list if args.contains in f]
if args.remove is not None:
file_list = [f for f in file_list if args.remove not in f]
if args.limit is not None:
limit = int(args.limit)
file_list = file_list[:limit]
global_state.data_paths = file_list
success_prompt('File path correctly loaded!')
def whitebox(args):
if global_state.target is None:
error_prompt('You have first to set a target.')
return
if args.type is None:
error_prompt('You have to set an attack type.')
error_prompt(f'Chose from this list: {BYTE_ATTACKS}')
return
if args.inject is None:
if args.type != PARTIAL_DOS:
error_prompt('You have to set an injection amount.')
return
else:
args.inject = 58
args.net = global_state.target
attack = create_correct_whitebox_attack(args)
global_state.attack = attack
success_prompt(f'Set up attack: {args.type}')
def predict(args):
if global_state.target is None:
error_prompt('First you need to set a target.')
return
if args.path is None:
if global_state.data_paths is None:
error_prompt('You have to give an input path.')
return
paths = global_state.data_paths
elif not os.path.isfile(args.path):
error_prompt(f'{args.path} does not exists.')
return
else:
paths = [args.path]
net = create_wrapper_for_global_target()
stats = {
'detected': 0,
'total': 0,
'confidence': 0,
}
for p in paths:
with open(p, 'rb') as handle:
code = handle.read()
info_prompt(f'Computing prediction for {p}')
code = CArray(np.frombuffer(code, dtype=np.uint8)).atleast_2d()
y_pred, confidence = net.predict(code, return_decision_function=True)
y_pred = y_pred.item()
score = confidence[0, 1].item()
stats['detected'] += int(y_pred != 0)
stats['total'] += 1
stats['confidence'] += score
info_prompt(f'predicted label: {y_pred}')
info_prompt(f'confidence: {score}')
print('-' * 20)
if stats['total'] >= 1:
separator_prompt()
success_prompt('Prediction stats:')
success_prompt(f'Detected: {stats["detected"]} / {stats["total"]}')
success_prompt(
f'Detection Rate: {stats["detected"] / stats["total"] * 100} %')
success_prompt(
f'Mean confidence: {stats["confidence"] / stats["total"]}')
def run(args):
if global_state.target is None:
error_prompt('You must first set a target to attack (`target` command).')
return
if global_state.attack is None:
error_prompt('You must first set an attack strategy (`whitebox` or `blackbox` commands).')
return
if global_state.data_paths is None:
error_prompt('You must first set which samples to use (`data` command).')
return
if not strategy_check():
return
if args.output is not None:
if not os.path.isdir(args.output):
os.mkdir(args.output)
success_prompt(f'Folder {args.output} created!')
if isinstance(global_state.attack, CEnd2EndMalwareEvasion):
whitebox_attack(args.output)
elif isinstance(global_state.attack, CBlackBoxProblem):
blackbox_attack(args.output)
def strategy_check():
if type(global_state.target) == CClassifierEmber and isinstance(global_state.attack, CClassifierEnd2EndMalware):
error_prompt('Can\'t use end-to-end attack against not differentiable model!')
return False
return True