<reponame>zhangbo1997/Malware-Detection
import glob
import numpy as np
import pandas as pd
import networkx as nx
import os
import re
from itertools import combinations_with_replacement
from itertools import combinations
import random
import concurrent.futures
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score
import scipy.sparse
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.ensemble import GradientBoostingClassifier
import json

def get_all_smali(all_smali):
    """
    It combines all smali files in all directionries of an APP into one big text string.

    Args:
        all_smali: a list of all smali files.
    """
    all_smali_in_one = ''
    for file in all_smali:
        with open(file) as f:
            content = f.read().strip('\n')

        all_smali_in_one += content + '\t0\n'
    all_smali_in_one = all_smali_in_one.replace('\n', '')
    return all_smali_in_one

def sample_malwares(num_to_sample):
    """
    Sample a certain number of malwares from the given directory.

    Args:
        num_to_sample: the number of malwares that needs to be sampled
        Returns: the list of malware directories.
    """
    malware_dirs = random.sample(glob.glob('/datasets/dsc180a-wi20-public/Malware/amd_data_smali/*/*/*'), num_to_sample)
    return malware_dirs

def true_pos_rate(y_true, y_pred):
    """
    Calculate the true postive (indicated by 1) rate of the prediction.
    Args:
        y_true: true label of data
        y_pred: prediction label of data
    """
    true_pos = np.sum(np.array(y_true)*np.array(y_pred))
    all_pos = sum(y_true)
    return true_pos/all_pos

def base_line_feature(app_dir):
    """
    Update dictionaries A, B, P when a new App comes in.

    Args:
        A: the app to api dictionary.
        B: the api to api dictionary in same package.
        P: the api to api dictionary in same package.
        app_dirs: a new app that comes in (name string). If it's malware,
                  then the app is a directory of the malware rather than name
    """

    ## this line get all smali files that is within this app directory
    all_smali = glob.glob(app_dir + "/**/*.smali", recursive = True)

    all_smali_in_one = get_all_smali(all_smali)

    codeblock_list = re.findall(r'\.method.*?\.end method', all_smali_in_one)

    num_code_block = len(codeblock_list)

    ## get all invoke calls from the file
    all_invoke_list = re.findall(r"(invoke-.*?->.*?)\s",all_smali_in_one)
    ## split the all_invoke_list into a list of tuples. Tuple format: (invoke method, api, package)
    all_invoke_tuples = [(invoke.split()[0], invoke.split()[-1], invoke.split()[-1].split('->')[0])
                         for invoke in all_invoke_list]

    num_unique_api = len(set([tuple_[1] for tuple_ in all_invoke_tuples]))
    num_unique_package = len(set([tuple_[2] for tuple_ in all_invoke_tuples]))

    num_api = len(all_invoke_list)
    return num_api, num_unique_api, num_code_block, num_unique_package



def baseline(num_benign, num_malware):
    benignware_dirs = glob.glob("./smali_files/*", recursive = True)
    benignware_dirs = random.sample(benignware_dirs, num_benign)
    malware_dirs = sample_malwares(num_malware)
    X_data = benignware_dirs + malware_dirs
    y_data = [0]*len(benignware_dirs)+[1]*len(malware_dirs)
    X_train, X_test, y_train, y_test = train_test_split(X_data, y_data, test_size=0.33, random_state=42)
    X_train_data = pd.Series(X_train).apply(lambda x: base_line_feature(x))
    X_test_data = pd.Series(X_test).apply(lambda x: base_line_feature(x))

    
    
    LR = LogisticRegression().fit([list(tuple_) for tuple_ in X_train_data], y_train)
    y_pred_LR = LR.predict([list(tuple_) for tuple_ in X_test_data])

    RF = RandomForestClassifier().fit([list(tuple_) for tuple_ in X_train_data], y_train)
    y_pred_RF = RF.predict([list(tuple_) for tuple_ in X_test_data])

    GBC = GradientBoostingClassifier().fit([list(tuple_) for tuple_ in X_train_data], y_train)
    y_pred_GBC = GBC.predict([list(tuple_) for tuple_ in X_test_data])

    acc = [
        accuracy_score(y_test, y_pred_LR),
        accuracy_score(y_test, y_pred_RF),
        accuracy_score(y_test, y_pred_GBC)
    ]

    tpr = [
        true_pos_rate(y_test, y_pred_LR),
        true_pos_rate(y_test, y_pred_RF),
        true_pos_rate(y_test, y_pred_GBC)
    ]
    df = pd.DataFrame(columns = ['Logistic Regression', 'Random Forest', 'Gradient Boosted Tree'], index = ['Acc.', 'True Positive Rate'])
    df.loc['Acc.',:] = acc
    df.loc['True Positive Rate',:] = tpr

    result_dir = json.load(open('../config/env.json'))['output-path']
    
    if not os.path.exists(result_dir):
        os.mkdir(result_dir)
    df.to_csv(result_dir+"/baseline_result.csv")
    print(df)

if __name__ == '__main__':
    baseline(200, 200)

import glob
import numpy as np
import pandas as pd
import networkx as nx
import os
import re
from itertools import combinations_with_replacement
from itertools import combinations
import random
import concurrent.futures
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MultiLabelBinarizer
from sklearn.svm import SVC
from sklearn.metrics import accuracy_score
import scipy.sparse
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.ensemble import GradientBoostingClassifier
import json

def get_all_smali(all_smali):
    """
    It combines all smali files in all directionries of an APP into one big text string.

    Args:
        all_smali: a list of all smali files.
    """
    all_smali_in_one = ''
    for file in all_smali:
        with open(file) as f:
            content = f.read().strip('\n')

        all_smali_in_one += content + '\t0\n'
    all_smali_in_one = all_smali_in_one.replace('\n', '')
    return all_smali_in_one

def sample_malwares(num_to_sample):
    """
    Sample a certain number of malwares from the given directory.

    Args:
        num_to_sample: the number of malwares that needs to be sampled
        Returns: the list of malware directories.
    """
    malware_dirs = random.sample(glob.glob('/datasets/dsc180a-wi20-public/Malware/amd_data_smali/*/*/*'), num_to_sample)
    return malware_dirs

def true_pos_rate(y_true, y_pred):
    """
    Calculate the true postive (indicated by 1) rate of the prediction.
    Args:
        y_true: true label of data
        y_pred: prediction label of data
    """
    true_pos = np.sum(np.array(y_true)*np.array(y_pred))
    all_pos = sum(y_true)
    return true_pos/all_pos

def base_line_feature(app_dir):
    """
    Update dictionaries A, B, P when a new App comes in.

    Args:
        A: the app to api dictionary.
        B: the api to api dictionary in same package.
        P: the api to api dictionary in same package.
        app_dirs: a new app that comes in (name string). If it's malware,
                  then the app is a directory of the malware rather than name
    """

    ## this line get all smali files that is within this app directory
    all_smali = glob.glob(app_dir + "/**/*.smali", recursive = True)

    all_smali_in_one = get_all_smali(all_smali)

    codeblock_list = re.findall(r'\.method.*?\.end method', all_smali_in_one)

    num_code_block = len(codeblock_list)

    ## get all invoke calls from the file
    all_invoke_list = re.findall(r"(invoke-.*?->.*?)\s",all_smali_in_one)
    ## split the all_invoke_list into a list of tuples. Tuple format: (invoke method, api, package)
    all_invoke_tuples = [(invoke.split()[0], invoke.split()[-1], invoke.split()[-1].split('->')[0])
                         for invoke in all_invoke_list]

    num_unique_api = len(set([tuple_[1] for tuple_ in all_invoke_tuples]))
    num_unique_package = len(set([tuple_[2] for tuple_ in all_invoke_tuples]))

    num_api = len(all_invoke_list)
    return num_api, num_unique_api, num_code_block, num_unique_package



def baseline(num_benign, num_malware):
    benignware_dirs = glob.glob("./smali_files/*", recursive = True)
    benignware_dirs = random.sample(benignware_dirs, num_benign)
    malware_dirs = sample_malwares(num_malware)
    X_data = benignware_dirs + malware_dirs
    y_data = [0]*len(benignware_dirs)+[1]*len(malware_dirs)
    X_train, X_test, y_train, y_test = train_test_split(X_data, y_data, test_size=0.33, random_state=42)
    X_train_data = pd.Series(X_train).apply(lambda x: base_line_feature(x))
    X_test_data = pd.Series(X_test).apply(lambda x: base_line_feature(x))

    
    
    LR = LogisticRegression().fit([list(tuple_) for tuple_ in X_train_data], y_train)
    y_pred_LR = LR.predict([list(tuple_) for tuple_ in X_test_data])

    RF = RandomForestClassifier().fit([list(tuple_) for tuple_ in X_train_data], y_train)
    y_pred_RF = RF.predict([list(tuple_) for tuple_ in X_test_data])

    GBC = GradientBoostingClassifier().fit([list(tuple_) for tuple_ in X_train_data], y_train)
    y_pred_GBC = GBC.predict([list(tuple_) for tuple_ in X_test_data])

    acc = [
        accuracy_score(y_test, y_pred_LR),
        accuracy_score(y_test, y_pred_RF),
        accuracy_score(y_test, y_pred_GBC)
    ]

    tpr = [
        true_pos_rate(y_test, y_pred_LR),
        true_pos_rate(y_test, y_pred_RF),
        true_pos_rate(y_test, y_pred_GBC)
    ]
    df = pd.DataFrame(columns = ['Logistic Regression', 'Random Forest', 'Gradient Boosted Tree'], index = ['Acc.', 'True Positive Rate'])
    df.loc['Acc.',:] = acc
    df.loc['True Positive Rate',:] = tpr

    result_dir = json.load(open('../config/env.json'))['output-path']
    
    if not os.path.exists(result_dir):
        os.mkdir(result_dir)
    df.to_csv(result_dir+"/baseline_result.csv")
    print(df)

if __name__ == '__main__':
    baseline(200, 200)
It combines all smali files in all directionries of an APP into one big text string.      Args:         all_smali: a list of all smali files. Sample a certain number of malwares from the given directory.      Args:         num_to_sample: the number of malwares that needs to be sampled         Returns: the list of malware directories. Calculate the true postive (indicated by 1) rate of the prediction.     Args:         y_true: true label of data         y_pred: prediction label of data Update dictionaries A, B, P when a new App comes in.      Args:         A: the app to api dictionary.         B: the api to api dictionary in same package.         P: the api to api dictionary in same package.         app_dirs: a new app that comes in (name string). If it's malware,                   then the app is a directory of the malware rather than name ## this line get all smali files that is within this app directory ## get all invoke calls from the file ## split the all_invoke_list into a list of tuples. Tuple format: (invoke method, api, package)