-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
47 lines (36 loc) · 1.62 KB
/
main.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
import argparse
import os
from static_analysis_apk import static_analysis, create_features
from gradient_boosting import create_model
from sklearn.model_selection import train_test_split
def label_apks(malware_dir, goodware_dir):
apks = []
labels = []
# Label malware APKs as 1
for apk_file in os.listdir(malware_dir):
apks.append(os.path.join(malware_dir, apk_file))
labels.append(1)
# Label benign APKs as 0
for apk_file in os.listdir(goodware_dir):
if apk_file.endswith('.apk') or apk_file.endswith('.zip'):
apks.append(os.path.join(goodware_dir, apk_file))
labels.append(0)
return apks, labels
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--malware', type=str, required=True, help='Path to malware directory')
parser.add_argument('--goodware', type=str, required=True, help='Path to benign directory')
args = parser.parse_args()
all_permissions = set()
all_api_calls = set()
apk_count = 0
apks, labels = label_apks(args.malware, args.goodware)
for apk in apks:
apk_count += 1
print(f"Gathering features for apk {str(apk_count)} out of {len(apk) - 1} ")
_, perms, api_calls = static_analysis(apk, all_permissions, all_api_calls)
all_permissions.update(perms)
all_api_calls.update(api_calls)
features = create_features(apks, all_permissions, all_api_calls)
X_train, X_test, y_train, y_test = train_test_split(features, labels, test_size=0.3, random_state=42)
clf = create_model(X_train, X_test, y_train, y_test, all_permissions, all_api_calls)