Malware_analysis_binary_RFF_CSM.py

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

start_time = time.time()

def run_script(sigma, dim, g_dim, m):
  cmd = ['python', 'script.py', '--sigma',  str(sigma), '--dim', str(dim), '--g_dim', str(g_dim), 'm', str(m)]
  subprocess.run(cmd)

# List of parameter sets for Hyperparameter optimization (alpha, subspace_dimension, gds_dimension, m_number_of_random_vectors) 
parameter_sets = [(alpha, i, 2*i + 28, 600) for alpha in [0.1, 0.3, 1, 5, 10] for i in range(30, 66)]


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 = []
m_rand = []
dimension = []
sgm = []


for sigma, dim, g_dim, m in parameter_sets:
  run_script(sigma, dim, g_dim, m)
  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','sigma', 'dim', 'train_x','train_y','test_x','test_y','g_dim' ,'m', 'Accuracy', 'm_rand', 'dimension', 'sgm', '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 scipy.linalg import eigh
  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




  def dual_vectors(K, n_subdims=None, higher=True, eps=1e-6):
      """
      Calc dual representation of vectors in kernel space
      Parameters:
      -----------
      K :  array-like, shape: (n_samples, n_samples)
          Grammian Matrix of X: K(X, X)
      n_subdims: int, default=None
          Number of vectors of dual vectors to return
      higher: boolean, default=None
          If True, this function returns eigenbasis corresponding to 
              higher `n_subdims` eigenvalues in descending order.
          If False, this function returns eigenbasis corresponding to 
              lower `n_subdims` eigenvalues in descending order.
      eps: float, default=1e-20
          lower limit of eigenvalues
      Returns:
      --------
      A : array-like, shape: (n_samples, n_samples)
          Dual replesentation vectors.
          it satisfies lambda[i] * A[i] @ A[i] == 1, where lambda[i] is i-th biggest eigenvalue
      e:  array-like, shape: (n_samples, )
          Eigen values descending sorted
      """

      eigvals = _get_eigvals(K.shape[0], n_subdims, higher)
      e, A = _eigen_basis(K, eigvals=eigvals)

      # replace if there are too small eigenvalues
      e[(e < eps)] = eps

      A = A / np.sqrt(e)

      return A, e

  def _prepare_X(X):
          """
          preprocessing data matricies X.
          normalize and transpose
          Parameters
          ----------
          X: list of 2d-arrays, (n_classes, n_samples, n_dims)
          """

          # transpose to make feature vectors as column vectors
          # this makes it easy to implement refering to formula
          X = [_X.T for _X in X]

          return X

  def _eigh(X, eigvals=None):
      """
      A wrapper function of numpy.linalg.eigh and scipy.linalg.eigh
      Parameters
      ----------
      X: array-like, shape (a, a)
          target symmetric matrix
      eigvals: tuple, (lo, hi)
          Indexes of the smallest and largest (in ascending order) eigenvalues and corresponding eigenvectors
          to be returned: 0 <= lo <= hi <= M-1. If omitted, all eigenvalues and eigenvectors are returned.
      Returns
      -------
      e: array-like, shape (a) or (n_dims)
          eigenvalues with descending order
      V: array-like, shape (a, a) or (a, n_dims)
          eigenvectors
      """

      if eigvals != None:
          e, V = eigh(X, eigvals=eigvals)
      else:
          # numpy's eigh is faster than scipy's when all calculating eigenvalues and eigenvectors
          e, V = np.linalg.eigh(X)

      e, V = e[::-1], V[:, ::-1]

      return e, V

  def _eigen_basis(X, eigvals=None):
      """
      Return subspace basis using PCA
      Parameters
      ----------
      X: array-like, shape (a, a)
          target matrix
      n_dims: integer
          number of basis
      Returns
      -------
      e: array-like, shape (a) or (n_dims)
          eigenvalues with descending order
      V: array-like, shape (a, a) or (a, n_dims)
          eigenvectors
      """

      try:
          e, V = _eigh(X, eigvals=eigvals)
      except np.linalg.LinAlgError:
          # if it not converges, try with tiny salt
          salt = 1e-8 * np.eye(X.shape[0])
          e, V = eigh(X + salt, eigvals=eigvals)

      return e, V

  def _get_eigvals(n, n_subdims, higher):
      """
      Culculate eigvals for eigh
      
      Parameters
      ----------
      n: int
      n_subdims: int, dimension of subspace
      higher: boolean, if True, use higher `n_subdim` basis
      Returns
      -------
      eigvals: tuple of 2 integers
      """

      if n_subdims is None:
          return None

      if higher:
          low = max(0, n - n_subdims)
          high = n - 1
      else:
          low = 0
          high = min(n - 1, n_subdims - 1)

      return low, high

  def subspace_bases(X, n_subdims=None, higher=True, return_eigvals=False):
      """
      Return subspace basis using PCA
      Parameters
      ----------
      X: array-like, shape (n_dimensions, n_vectors)
          data matrix
      n_subdims: integer
          number of subspace dimension
      higher: bool
          if True, this function returns eigenvectors collesponding
          top-`n_subdims` eigenvalues. default is True.
      return_eigvals: bool
          if True, this function also returns eigenvalues.
      Returns
      -------
      V: array-like, shape (n_dimensions, n_subdims)
          bases matrix
      w: array-like shape (n_subdims)
          eigenvalues
      """

      if X.shape[0] <= X.shape[1]:
          eigvals = _get_eigvals(X.shape[0], n_subdims, higher)
          # get eigenvectors of autocorrelation matrix X @ X.T
          w, V = _eigen_basis(X @ X.T, eigvals=eigvals)
      else:
          # use linear kernel to get eigenvectors
          A, w = dual_vectors(X.T @ X, n_subdims=n_subdims, higher=higher)
          V = X @ A

      if return_eigvals:
          return V, w
      else:
          return V

  def _gds_projection(gds, bases):
          """
          GDS projection.
          Projected bases will be normalized and orthogonalized.
          Parameters
          ----------
          bases: arrays, (n_dims, n_subdims)
          Returns
          -------
          bases: arrays, (n_gds_dims, n_subdims)
          """

          # bases_proj, (n_gds_dims, n_subdims)
          bases_proj = np.matmul(gds.T, bases)
          qr = np.vectorize(np.linalg.qr, signature="(n,m)->(n,m),(m,m)")
          bases, _ = qr(bases_proj)
          return bases


  def z(X, w, b, m):
      #print(w.shape, X.shape, b.shape)
      return np.sqrt(2 / m) * np.cos((w @ X).T + b).T

  def get_newX(X,dim,g_dim,m,sigma):
      n_gds_dims = g_dim 
      n_subdims = dim
      dic = None
      dic = []

      n_dims = X[0].shape[0]
      if sigma is None:
          sigma = np.sqrt(n_dims / 2)

      
      newX = [] 
      for _X in X: 
          _newX = []
          for n in range(n_approx):
              _newX.append(z(_X, w7[n], b7[n], m))
          _newX = np.stack(_newX, axis=0).mean(axis=0)
          newX.append(_newX)

      newX_list_of_list = []
      for x in newX:
          newX_list_of_list.append([x])

      for newXx in newX_list_of_list:
          newXx = np.array(newXx)
          dic_ = [subspace_bases(_X, n_subdims) for _X in newXx] 
          dic.append(dic_)
          
      dic = np.array(dic)
      dic_ = dic[:, 0, :, :] 

      all_bases = np.hstack((dic_[0], dic_[1])) 
        
      if 0.0 < n_gds_dims <= 1.0:
          n_gds_dims = int(all_bases.shape[1] * n_gds_dims)
      else:
          n_gds_dims = n_gds_dims

      gds = subspace_bases(all_bases, n_gds_dims, higher=False)
      dic = _gds_projection(gds,dic_)
      return dic,gds
  # #----------------------------------------------------------------------------------------

  print("sigma:",sigma,"subdim:",dim,"gdsdim:",g_dim,"m_random",m)   



  #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)
  #----------------------------------------------------------------------------------------




  n_dims = train_x[0].shape[1]
  n_approx = 1

  w = []
  b = []
  for n in range(n_approx):
      w.append(np.random.randn(m, n_dims) / sigma)
      b.append(np.random.rand(m) * 2 * np.pi)
  w7 = np.stack(w)
  b7 = np.stack(b)
  # print(w7.shape, b7.shape)


  train_xx = _prepare_X(train_x)



  # rff for training data + PCA 
  New_X,gds = get_newX(train_xx,dim,g_dim,m,sigma)
  #------------------------



  #calculate similarity  
  similarity_all = []
  for i in range(len(test_x)):
      data_input = test_x[i]
      #print(w7.shape, b7.shape)
      z_in = z(data_input.T, w7[0,:], b7[0, :], m)
      z_projected = gds.T @ z_in
      #print(z_projected.shape)
      #print(data_input.shape)
      #print(z_in.shape)
      similarity_one = []
      for subspace_class_i in New_X:
          # print(subspace_class_i.shape, z_projected.shape)
          length = np.linalg.norm(subspace_class_i.T @ z_projected, ord = 2)
          # print(length)
          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))



  Accuracy.append(accuracy_score(similarity_all, test_y))
  m_rand.append(m)
  dimension.append(dim)
  sgm.append(sigma)

  conf_mat = confusion_matrix(test_y, similarity_all)
  sns.heatmap(conf_mat,annot=True,fmt='d',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")