Malware_analysis_binary_KCSM.py

#  %%
import subprocess
import hdf5storage
import numpy as np
import time

start_time = time.time()

def run_script(sgm, subdim, gdsdim):
  cmd = ['python', 'script.py', '--sgm', str(sgm), '--subdim', str(subdim), '--gdsdim', str(gdsdim)]
  subprocess.run(cmd)

# List of parameter sets for Hyperparameter optimization (alpha, subspace_dimension, gds_dimension) 
parameter_sets = [(0.1, i, 2*i + 28) for i in range(30, 101)]


X_safe = hdf5storage.loadmat('X_safe.mat')
X_mal = hdf5storage.loadmat('X_mal.mat')
in_mal = hdf5storage.loadmat('in_mal.mat')

X_safe = np.array(X_safe['X_safe']).T
X_mal = np.array(X_mal['X_mal']).T
in_mal = np.array(in_mal['in_mal']).T

rows, cols = X_safe.shape
# take 70% for safe train 
train_rows = int(0.7 * rows)
X_safe_chosen = X_safe[:train_rows, :]

#-------------------------------------------------------------------------------
X_mal_chosen = X_mal #take everything without fold splitting
#-------------------------------------------------------------------------------


train_x = [X_safe_chosen, X_mal_chosen] 
train_y = np.array([0, 1])
# Assign the remaining 30% for safe test
in_safe_chosen = X_safe[train_rows:, :]
#-------------------------------------------------------------------------------
in_mal_chosen = in_mal #take everything without fold splitting
#-------------------------------------------------------------------------------


test_x = np.concatenate([in_safe_chosen, in_mal_chosen]) 
test_x = [test_x[i, np.newaxis, :] for i in range(test_x.shape[0])]
test_y0 = np.full(len(in_safe_chosen), 0)
test_y1 = np.full(len(in_mal_chosen), 1)
test_y = np.concatenate([test_y0, test_y1]) 
Accuracy = []
dimension_safe = []
dimension_mal = []




for sgm, subdim, gdsdim in parameter_sets:
  run_script(sgm, subdim, gdsdim)
  for var in list(globals().keys()):
    if var not in ['function variables','accuracy_score','gauss_class_mat_diff','gauss_gram_mat','gauss_gram_two','gauss_projection_diff','get_ipython','jit','parentpath1','rand','randint','run_script','trace_this_thread','train_test_split','special variables','__','___','__doc__','__file__','__loader__','__name__','__package__','__spec__','__vsc_ipynb_file__','__builtin__','__builtins__','debugpy','exit','hdf5storage','np','os','pd','quit','random','subprocess','sys','_VSCODE_hashlib','_VSCODE_types','_dh','_VSCODE_compute_hash','_VSCODE_wrapped_run_cell','normalize','sgm', 'subdim', 'train_x','train_y','test_x','test_y','gdsdim', 'time', 'start_time']:
      del globals()[var]

  import os
  import numpy as np
  from numpy.random import randint, rand
  from sklearn.model_selection import train_test_split
  from sklearn.metrics import accuracy_score
  import hdf5storage
  import random
  import pandas as pd
  from numpy.random import randint, rand
  from numba import jit, void, f8, njit
  from sklearn.preprocessing import normalize
  import matplotlib.pyplot as plt
  from sklearn.metrics import confusion_matrix
  from sklearn.metrics import precision_score, recall_score, f1_score
  import seaborn as sns

  @jit(void(f8[:, :], f8[:, :]))
  def gauss_gram_mat(x, K):
    n_points = len(x)
    n_dim = len(x[0])
    b = 0

    for j in range(n_points):
      for i in range(n_points):
        for k in range(n_dim):
          b = (x[i][k] - x[j][k])
          K[i][j] += b * b

  @jit
  def gauss_class_mat_diff(K_d):
      for c1 in range(class_num):
          for c2 in range(class_num):
              for i1 in range(subdim):
                  for j2 in range(subdim):
                      for s1 in range(class_info[c1]):
                          for t2 in range(class_info[c2]):
                              K_d[subdim * c1 + i1, subdim * c2 + j2] += K_cl[class_index[c1] + s1, i1] * K_cl[class_index[c2] + t2, j2] * K_all[class_index[c1] + s1, class_index[c2] + t2]

  @jit
  def gauss_projection_diff(x, K_p):
      for i_data in range(x.shape[0]):
          for i_gds in range(b.shape[1]):
              for c1 in range(class_num):
                  for i1 in range(subdim):
                      for s1 in range(class_info[c1]):
                          K_p[i_gds, i_data] += b[i_gds, subdim * c1 + i1] * K_cl[class_index[c1] + s1, i1] * x[i_data, class_index[c1] + s1]

  @jit
  def gauss_gram_two(X, Y, K_t):
      for i in range(X.shape[0]):
          for j in range(Y.shape[0]):
              for k in range(X.shape[1]):
                  b = (X[i][k] - Y[j][k])
                  K_t[i][j] += b * b


  def dual_vectors(K, n_subdims=None, higher=True, elim=False, eps=1e-6):
      e, A = np.linalg.eigh(K)
      e[(e < eps)] = eps

      A = A / np.sqrt(e)

      if elim:
          e = e[(e > eps)]
          A = A[:, (e > eps)]

      if higher:
          return A[:, -n_subdims:]

      return A[:, :n_subdims]

  print("sigma:",sgm,"subdim:",subdim,"gdsdim:",gdsdim)    

  #parameter-------------------------------------------------------------------------------
  class_num = 2                               
  class_info = np.array([1750, 8404]) ##for Safe + Malimg dataset
  #----------------------------------------------------------------------------------------


  #normalization---------------------------------------------------------------------------
  class_index = []
  count_c = 0
  for class_i in class_info:
      if count_c == 0:
          class_index.append(0)
          class_index.append(class_info[0])
      else:
          class_index.append(class_index[count_c] + class_info[count_c])
      count_c += 1
  class_index = np.array(class_index)



  K_class = []
  count_class = 0
  for train_data in train_x:
      K = np.zeros((train_data.shape[0], train_data.shape[0]))
      gauss_gram_mat(train_data, K)
      K = np.exp(- K / (2 * sgm))
      X1_coeffs = dual_vectors(K, n_subdims=subdim)
      K_class.append(X1_coeffs)

  count_k = 0
  for i in range(class_num):
      if i == 0:
          K_cl = K_class[i]
      else:
          K_cl = np.concatenate([K_cl, K_class[i]])     

  K_D = np.zeros((subdim * class_num, subdim * class_num)) 

  for i in range(class_num):
      if i == 0:
          train_all = train_x[i]
      else:
          train_all = np.concatenate([train_all, train_x[i]]) 

  K_all = np.zeros((train_all.shape[0], train_all.shape[0])) 
                                              
  gauss_gram_mat(train_all, K_all) 
  K_all = np.exp(- K_all / (2 * sgm)) 

  gauss_class_mat_diff(K_D) 
  B, b = np.linalg.eigh(K_D)
  b = b[:, 0:gdsdim]
  print("training complete")
  #-----------------------------------------------------------------------------------------


  #project data onto gds--------------------------------------------------------------------
  test_np = np.array(test_x)
  test_np = test_np.reshape([test_np.shape[0], test_np.shape[2]])

  train_data_projected = np.zeros((gdsdim, train_all.shape[0]))
  gauss_projection_diff(K_all, train_data_projected)
  print("projection complete 1 (train(class) data)")

  K_two = np.zeros((test_np.shape[0], train_all.shape[0]))
  gauss_gram_two(test_np, train_all, K_two)
  K_two = np.exp(- K_two / (2 * sgm))

  test_data_projected = np.zeros((gdsdim, test_np.shape[0]))
  gauss_projection_diff(K_two, test_data_projected)
  print("projection complete 2 (test data)")
  #-----------------------------------------------------------------------------------------
  

#generate linear subspace for train data projected onto gds-------------------------------
  subspace_class = []
  for i in range(class_num):
      data_projected = train_data_projected[:, class_index[i]:class_index[i+1]]
      w, v = np.linalg.eigh(data_projected @ data_projected.T)
      w, v = w[::-1], v[:, ::-1]
      rank = np.linalg.matrix_rank(K)
      w, v = w[:rank], v[:, :rank]
      d = v[:, 0:subdim]
      subspace_class.append(d)
  #-----------------------------------------------------------------------------------------


  #calculate projection length--------------------------------------------------------------
  similarity_all = []
  for i in range(test_np.shape[0]):
      data_input = test_data_projected[:, i]
      similarity_one = []
      for subspace_class_i in subspace_class:
          length = np.linalg.norm(subspace_class_i.T @ data_input, ord = 2)
          similarity_one.append(length)
      similarity_all.append(np.argmax(np.array(similarity_one))) 
  #------------------------------------------------------------------------------------------

  precision = precision_score(test_y, similarity_all)
  recall = recall_score(test_y, similarity_all)
  f1 = f1_score(test_y, similarity_all)
  print("Precision: {:.2f}".format(precision))
  print("Recall: {:.2f}".format(recall))
  print("F1 Score: {:.2f}".format(f1))
  print("accuracy:", accuracy_score(similarity_all, test_y))



  conf_mat = confusion_matrix(test_y, similarity_all)
  precision_class0 = conf_mat[0, 0] / (conf_mat[0, 0] + conf_mat[1, 0])
  precision_class1 = conf_mat[1, 1] / (conf_mat[1, 1] + conf_mat[0, 1])
  conf_mat_precision = np.array([[precision_class0, 1 - precision_class0], [1 - precision_class1, precision_class1]])
  sns.heatmap(conf_mat_precision, annot=True, fmt='.2f', cmap='OrRd', xticklabels=['Safe', 'Malware'], yticklabels=['Safe', 'Malware'])
  plt.xlabel('Predicted labels')
  plt.ylabel('True labels')
  plt.tight_layout()
  plt.show()

  end_time = time.time()
  elapsed_time = end_time - start_time
  print("Elapsed time: ", elapsed_time, " seconds")