def generate_mutated_malware(file, model, args):
pe = pefeatures2.PEFeatureExtractor2()
rn = RangeNormalize(-0.5,0.5)
info("[*] Reading file : " + str(file))
bytez = []
with open(str(file), 'rb') as infile:
bytez = infile.read()
for t in track(range(1, args.rl_mutations) , description="Generating mutation ...", transient=True):
state = pe.extract( bytez )
state_norm = rn(state)
state_norm = torch.from_numpy(state_norm).float().unsqueeze(0).to(device)
actions = model.forward(state_norm)
action = torch.argmax(actions).item()
action = ACTION_LOOKUP[action]
bytez = bytes(manipulate.modify_without_breaking(bytez, [action]))
new_score = interface.get_score_local( bytez )
if(new_score < interface.local_model_threshold):
break
if not os.path.exists(args.o):
os.mkdir(args.o)
info("[*] output directory has been created at : " + str(args.o))
output_file = os.path.join(args.o, "mutated_" + str(os.path.basename(file)))
info("[*] Writing mutated file to : " + str(output_file) + "\n\n")
with open(str(output_file), mode='wb') as file1:
file1.write(bytes(bytez))
return
def _take_action(self, action_index):
assert action_index < len(ACTION_LOOKUP)
action = ACTION_LOOKUP[action_index]
print(action)
self.history[self.sha256]['actions'].append(action)
self.bytez = bytes(
manipulate.modify_without_breaking(self.bytez, [action]))
def evaluate(action_function):
success = []
misclassified = []
for sha256 in sha256_holdout:
print("#########{}#########".format(sha256))
success_dict = defaultdict(list)
bytez_o = interface.fetch_file(sha256)
label = interface.get_label_windefender(bytez_o)
if label == 0.0:
misclassified.append(sha256)
continue # already misclassified, move along
bytez = bytez_o
for _ in range(MAXTURNS):
action = action_function(bytez)
print(action)
success_dict[sha256].append(action)
bytez = manipulate.modify_without_breaking(bytez, [action])
new_label = interface.get_label_windefender(bytez)
if new_label == 0.0:
success.append(success_dict[sha256])
os.mkdir(PATH_SAVE + '/' + sha256)
with open(PATH_SAVE + '/' + sha256 + '/orign', 'wb') as f:
f.write(bytez_o)
with open(PATH_SAVE + '/' + sha256 + '/modified', 'wb') as f:
f.write(bytez)
break
return success, misclassified # evasion accuracy is len(success) / len(sha256_holdout)
def _take_action(self, action_index):
assert action_index < len(ACTION_LOOKUP)
action = ACTION_LOOKUP[action_index]
print(action)
self.history[self.sha256]['actions'].append(action)
self.bytez = bytes(
manipulate.modify_without_breaking(self.bytez, [action]))
def test_model():
T = 80 # total mutations allowed
success = 0
rn = dqeaf.RangeNormalize(-0.5, 0.5)
fe = pefeatures.PEFeatureExtractor()
episode = 0
for file in onlyfiles:
try:
with open(os.path.join(input_folder, file), 'rb') as infile:
bytez = infile.read()
except IOError:
raise FileRetrievalFailure("Unable to read sha256 from")
state = fe.extract(bytez)
state_norm = rn(state)
episode = episode + 1
state_norm = torch.from_numpy(state_norm).float().unsqueeze(0).to(
device)
for mutation in range(1, T):
actions = model.forward(state_norm)
print(actions)
action = torch.argmax(actions).item()
action = ACTION_LOOKUP[action]
bytez = manipulate.modify_without_breaking(bytez, [action])
new_label = interface.get_score_local(bytez)
print('episode : ' + str(episode))
print('mutation : ' + str(mutation))
print('test action : ' + str(action))
print('new label : ' + str(new_label))
state = fe.extract(bytez)
state_norm = rn(state)
state_norm = torch.from_numpy(state_norm).float().unsqueeze(0).to(
device)
if (new_label < 0.90):
with open(os.path.join(output_folder, file + '.exe'),
mode='wb') as file1:
file1.write(bytes(bytez))
break
def evaluate(action_function):
success = []
misclassified = []
for sha256 in sha256_holdout:
success_dict = defaultdict(list)
bytez = interface.fetch_file(sha256)
label = interface.get_label_local(bytez)
if label == 0.0:
misclassified.append(sha256)
continue # already misclassified, move along
for _ in range(MAXTURNS):
action = action_function(bytez)
print(action)
success_dict[sha256].append(action)
bytez = manipulate.modify_without_breaking(bytez, [action])
new_label = interface.get_label_local(bytez)
if new_label == 0.0:
success.append(success_dict)
break
return success, misclassified # evasion accuracy is len(success) / len(sha256_holdout)
def evaluate( action_function ):
success=[]
misclassified = []
for sha256 in sha256_holdout:
success_dict = defaultdict(list)
bytez = interface.fetch_file(sha256)
label = interface.get_label_local(bytez)
if label == 0.0:
misclassified.append(sha256)
continue # already misclassified, move along
for _ in range(MAXTURNS):
action = action_function( bytez )
print(action)
success_dict[sha256].append(action)
bytez = manipulate.modify_without_breaking( bytez, [action] )
new_label = interface.get_label_local( bytez )
if new_label == 0.0:
success.append(success_dict)
break
return success, misclassified # evasion accuracy is len(success) / len(sha256_holdout)
# 本程序用于逐个运行lief修改PE文件的action，来删除不适合实验的PE文件
file_list = interface.get_available_sha256()
action_test_list = {
'overlay_append',
'imports_append_org',
'section_add',
'remove_signature'
}
# run the tests
index = 1
for sha256 in file_list:
print("{}:原文件:{}".format(index, sha256))
try:
for action_name in action_test_list:
print("使用{}动作来修改".format(action_name))
action = manipulate.ACTION_TABLE[action_name]
print("action{}".format(action))
bytez = interface.fetch_file(sha256)
bytez = bytes(manipulate.modify_without_breaking(bytez, [action]))
with open(sha256 + "_" + action_name, 'wb') as outfile:
outfile.write(bytez)
except Exception as e:
print("e{}".format(e))
index += 1
# sort and print result list
print("total:{}".format(file_list.__len__()))
import hashlib
from gym_malware.envs.controls import manipulate2 as manipulate
from gym_malware.envs.utils import interface
# np.random.seed(322333)
# random_action = lambda bytez: np.random.choice( list(manipulate.ACTION_TABLE.keys()) )
from other import lief_test
random_action = lambda bytez: 'section_append'
# original 32bit putty.exe in sha256 file name
fileName = 'VirusShare_0a7e4d6c6006ed4a33b5fe0c181062c0'
bytez = interface.fetch_file(fileName)
lief_test.showPE(fileName)
action = random_action(bytez)
print(action)
bytez_mod = manipulate.modify_without_breaking(bytez, [action])
with open("putty-mod32.exe", "wb") as new_file:
new_file.write(bytez_mod)
lief_test.showPE("putty-mod32.exe")
