pefile 악성코드 판단하는 CNN 모델 - 데이터 전처리

 PE file의 구조 중 .text부분을 캡스톤을 이용한 opcode시퀀스 추출 후 CountVectorizer로 벡터화 후 모델 테스트시 생각보다 성능이 굉장히 좋지 않았다 다른방법을 생각 중  .text의 binary를 추출하여 image화 한 후 CNN 모델로 테스트 해보기로 하였다

**pe file의 정적 피처들의 단점은 조작을 가하거나 패킹이 되있을 시 정상적으로 데이터가 추출이 되지 않는 단점이 있다.

 

image로 할 pefile의 .text 부분

1. 첫번째로 pefile을 다룰 python 라이브러리인 pefile을 설치

pip install pefile

 

2. pefile을 dir 내장함수로 어떠한 method들이 있는지 간단하게 확인

pe = pefile.PE(file_path)
pe
<pefile.PE at 0x7fa10edab4c0>

dir(pe)
['DIRECTORY_ENTRY_IMPORT',
 'DOS_HEADER',
 'FILE_HEADER',
 'FileAlignment_Warning',
 'NT_HEADERS',
 'OPTIONAL_HEADER',
 'PE_TYPE',
 'RICH_HEADER',
 'SectionAlignment_Warning',

목표인 sections가 확인

 sections을 print

for section in pe.sections:
    print(section)

 

notepad.exe의 .text부분을 PointerToRawData가 가리키는 것을 확인

시작 지점과 size를 확인 후에 해당 부분을 긁어서 numpy array로 가져온다.

def opcode_Get(file_path):
    try:
        pe = pefile.PE(file_path,fast_load=True)

        for section in pe.sections:
            if '.text' in str(section.Name):
                entry = section.PointerToRawData - 1
                end = section.SizeOfRawData + entry
                raw_data = pe.__data__[entry:end]
                data = np.frombuffer(raw_data, dtype = np.float32)                
        #return raw_data
        return np.nan_to_num(data)            
    except: 
        return
opcode_Get(file_path)
-> array([ 2.695370e+36, -9.690833e+24,  5.687310e-41, ...,  0.000000e+00,
        0.000000e+00,  0.000000e+00], dtype=float32)

생각했던 첫번째 방법은

byte array를 가져 온 후 -> image 화 -> pillow를 이용한 resize -> numpy array화 -> CNN 모델 학습 이였지만,

byte array -> numpy array로 변환 후 numpy array에서 64*64 = 4096으로 배열을 마추도록 하였다.

(차 후에 첫번째 방법 진행 예정)

res = []
for x in tqdm(file['raw_data']):
    s = int(x.shape[0]*0.3)
    e = int(x.shape[0]*0.9)
    array = x[s:e].copy()
    array.resize(length*length)
    res.append(array.reshape(length,length))
np_arr = np.stack([x for x in res])

numpy array들을 file이라는 dataframe의 raw_data column에 담은 후에

raw_data의 30% ~ 90% -> 중간 영역만 짤라서

가져와 length 길이만큼 resize 후 reshape로 2차원 numpy array로 변환 작업

(raw_data의 전체 byte array를 사용하는 건 비 효율적이라고 생각이되서 임의로 짜른 부분 차 후 모델 성능 테스트때 100%와 60%의 성능 테스트 예정)

 

**데이터 전처리 작업 완료

 

총 코드

import pefile
import numpy as np
import os
from tqdm import tqdm
import pandas as pd

PATH = '파일들 폴더'
filename = os.listdir(PATH)

def opcode_Get(file_path):
    try:
        pe = pefile.PE(file_path,fast_load=True)

        for section in pe.sections:
            if '.text' in str(section.Name):
                entry = section.PointerToRawData - 1
                end = section.SizeOfRawData + entry
                raw_data = pe.__data__[entry:end]
                data = np.frombuffer(raw_data, dtype = np.float32)                
        #return raw_data
        return np.nan_to_num(data)
            
    except: 
        return
        
res = []
for x in tqdm(filename):
    target = f'{PATH}{x}'
    if os.path.isfile(target):        
        output = opcode_Get(target)
        if output is not None:
            res.append((x,output))
            
df = pd.DataFrame(res,columns = ['fid','raw_data']) #dataframe으로 데이터 확인을 위해 만들었지만 사용하지 않아도 된다.

length = np.ceil(np.sqrt(merge['data'].apply(lambda x:len(x)).mean()*0.6)).astype(int)

length = 64
arr = []
for x in tqdm(merge['data']):
    s = int(x.shape[0]*0.3)
    e = int(x.shape[0]*0.9)
    x_cp = x[s:e].copy()
    x_cp.resize(length*length)
    arr.append(x_cp.reshape(length,length))
train_data = np.stack([x for x in res])