-
Notifications
You must be signed in to change notification settings - Fork 93
/
pe_header_features.py
119 lines (101 loc) · 5.53 KB
/
pe_header_features.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
115
116
117
118
119
"""Extract LIEF features from PE files"""
from h2oaicore.transformer_utils import CustomTransformer
import datatable as dt
import numpy as np
class PEHeaderFeatures(CustomTransformer):
_unsupervised = True
_modules_needed_by_name = ['lief==0.14.1']
_regression = True
_binary = True
_multiclass = True
_is_reproducible = True
_parallel_task = True # if enabled, params_base['n_jobs'] will be >= 1 (adaptive to system), otherwise 1
_can_use_gpu = True # if enabled, will use special job scheduler for GPUs
_can_use_multi_gpu = True # if enabled, can get access to multiple GPUs for single transformer (experimental)
_numeric_output = True
@staticmethod
def get_default_properties():
return dict(col_type="text", min_cols=1, max_cols=1, relative_importance=1)
@staticmethod
def do_acceptance_test():
return False
def fit_transform(self, X: dt.Frame, y: np.array = None):
return self.transform(X)
def load_pe(self, file_path):
with open(file_path, 'rb') as f:
bytez = bytearray(f.read())
return (bytez)
def header_features(self, lief_binary):
from sklearn.feature_extraction import FeatureHasher
header = lief_binary.header
opt_header = lief_binary.optional_header
features = {}
# Header features
features['Header_time_date_stamps'] = header.time_date_stamps
features['Header_sizeof_optional_header'] = header.sizeof_optional_header
for i, x in enumerate(FeatureHasher(10, input_type='string').transform([str(header.machine)]).toarray()[0]):
features.update({f'Header_machine_hash_{i}': x})
for i, x in enumerate(FeatureHasher(10, input_type='string').transform(
[str(c) for c in header.characteristics_list]).toarray()[0]):
features.update({f'Header_characteristics_hash_{i}': x})
# Optional Header features
for i, x in enumerate(
FeatureHasher(10, input_type='string').transform([str(opt_header.subsystem)]).toarray()[0]):
features.update({f'Optional_Header_subsystem_hash_{i}': x})
for i, x in enumerate(FeatureHasher(10, input_type='string').transform(
[str(c) for c in opt_header.dll_characteristics_lists]).toarray()[0]):
features.update({f'Optional_Header_dll_characteristics_hash_{i}': x})
for i, x in enumerate(FeatureHasher(10, input_type='string').transform([str(opt_header.magic)]).toarray()[0]):
features.update({f'Optional_Header_magic_hash_{i}': x})
features['major_image_version'] = opt_header.major_image_version
features['minor_image_version'] = opt_header.minor_image_version
features['major_linker_version'] = opt_header.major_linker_version
features['minor_linker_version'] = opt_header.minor_linker_version
features[
'major_operating_system_version'] = opt_header.major_operating_system_version
features[
'minor_operating_system_version'] = opt_header.minor_operating_system_version
features['major_subsystem_version'] = opt_header.major_subsystem_version
features['minor_subsystem_version'] = opt_header.minor_subsystem_version
features['sizeof_code'] = opt_header.sizeof_code
features['sizeof_headers'] = opt_header.sizeof_headers
features['sizeof_heap_commit'] = opt_header.sizeof_heap_commit
return features
def get_header_features(self, file_path):
import lief
try:
pe_bytez = self.load_pe(file_path)
lief_binary = lief.PE.parse(list(pe_bytez))
X = self.header_features(lief_binary)
return X
except:
X = {'Header_time_date_stamps': 0,
'Header_sizeof_optional_header': 0}
X.update({f'Header_machine_hash_{i}': 0 for i in range(10)})
X.update({f'Header_characteristics_hash_{i}': 0 for i in range(10)})
X.update({f'Optional_Header_subsystem_hash_{i}': 0 for i in range(10)})
X.update({f'Optional_Header_dll_characteristics_hash_{i}': 0 for i in range(10)})
X.update({f'Optional_Header_magic_hash_{i}': 0 for i in range(10)})
X.update({f'{feature_name}': 0 for feature_name in ['major_image_version',
'minor_image_version',
'major_linker_version',
'minor_linker_version',
'major_operating_system_version',
'minor_operating_system_version',
'major_subsystem_version',
'minor_subsystem_version',
'sizeof_code',
'sizeof_headers',
'sizeof_heap_commit']})
return X
def transform(self, X: dt.Frame):
import pandas as pd
ret_df = pd.DataFrame(
[
self.get_header_features(x)
for x in X.to_pandas().values[:, 0]
]
)
self._output_feature_names = ret_df.columns.to_list()
self._feature_desc = self._output_feature_names
return ret_df