-
Notifications
You must be signed in to change notification settings - Fork 0
/
shifter.py
114 lines (92 loc) · 3.3 KB
/
shifter.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
import io
import itertools
import numpy as np
import os
import pefile
import random
class PEShifter:
"""
This class receives a PE file as input and operates on it by reordering
its sections.
"""
def __init__(self,
_input,
output_path=None,
verbose=False):
# Check if input is a path to a file or raw data
if isinstance(_input, str) and os.path.isfile(_input):
self.pe = pefile.PE(name=_input, fast_load=True)
self.input_buffer = open(_input, 'rb')
elif isinstance(_input, np.ndarray):
self.pe = pefile.PE(data=_input.tobytes())
self.input_buffer = io.BytesIO(_input)
else:
raise ValueError(f"Input has unknown type {type(_input)}.")
self.output_path = output_path
self.output_buffer = io.BytesIO()
self.verbose = verbose
return
def combine(self, offsets, data):
"""
Returns combination of two lists without repetition (kind of a cartesian product)
https://www.adamsmith.haus/python/answers/how-to-get-all-unique-combinations-of-two-lists-in-python
Args:
offsets (list): List with PointerToRawData offsets
data (bytes): Section data
Returns:
list: List of combinations
"""
all_combinations = []
list1_permutations = itertools.permutations(offsets, len(data))
for each_permutation in list1_permutations:
zipped = zip(each_permutation, data)
all_combinations.append(list(zipped))
return all_combinations
def run(self):
# Copy raw data to output
self.input_buffer.seek(0)
self.output_buffer.write(self.input_buffer.read())
self.input_buffer.seek(0)
self.output_buffer.seek(0)
# Read sections info
orig = []
offsets = []
data = []
i = 0
for s in self.pe.sections:
if s.PointerToRawData > 0:
orig.append(i)
# Tuple with offset where PointerToRawData is written at and its actual value
offsets.append([s.get_file_offset() + 20, s.PointerToRawData])
data.append(s.get_data())
i += 1
combinations = self.combine(offsets, data)
# We randomly get the combination
n_valid_combinations = len(combinations)
combination_idx = random.randint(0, n_valid_combinations-1)
chosen_combination = combinations[combination_idx]
if self.verbose:
print(f"\n\t[*] File has {len(n_valid_combinations)} combinations. {combination_idx} idx was chosen;")
final_data = None
for o, d in chosen_combination:
# Update PointerToRawData
self.output_buffer.seek(o[0])
self.output_buffer.write(o[1].to_bytes(4, 'little'))
# Update data
self.output_buffer.seek(o[1])
# self.output_buffer.write(os.urandom(8192))
self.output_buffer.write(d)
self.output_buffer.seek(0)
final_data = self.output_buffer.read()
# Write output path
if self.output_path is not None:
out_file_path = f"{self.output_path}_comb_{i}.exe"
out_file = open(out_file_path, 'wb')
out_file.write(final_data)
if self.verbose:
pe = pefile.PE(data=final_data)
print(f"{len(final_data)} bytes written to => {out_file_path}")
print(f"Final data: {pe.sections}")
self.input_buffer.close()
self.output_buffer.close()
return final_data