Static-Malware-Detection/identify_exe.py at master · atharv67/Static-Malware-Detection · GitHub

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
import re
import sys
import pickle
import pandas as pd
import tensorflow as tf
from extract_info import extract_info
from check_malware import check_malware


def get_data():
    '''
        Extracting data from application
    '''
    no_of_arg = len(sys.argv)-1
    arg = sys.argv[1:]

    for i in range(no_of_arg):
        if((arg[i][0]=="'") or (arg[i][0]=='"')):
            arg[i] = arg[i][1:-1]

    tupled_data = []
    for i in arg:
        result = extract_info(i)
        tupled_data.append(list(result.values()))
    test_df = pd.DataFrame(data=tupled_data, columns=list(result.keys()))

    classes, prob, pred = check_malware(test_df)
    print('--------------------------------------')
    print('Application => Category => Probability')
    print('--------------------------------------')
    for i in range(test_df.shape[0]):
        # for ML model
        # print(app[i], ' => ', classes[i], ' => ', prob[i][pred[i]])

        # for DL model
        print(arg[i], ' => ', classes[i], ' => ', prob[i][0])


if __name__ == "__main__":
    get_data()