수야나도 블로그

이 포스트는 프로젝트 진행 중 발생한 모든 오류 해결 과정을 포함하고 있습니다.
다소 긴 내용이 될 수 있으므로, 필요한 부분만 참고하시기 바랍니다.

1. 레이블링 이란

레이블링 = 라벨링은 학습할 데이터에 정답을 표시하는 것으로, 있다면 이것은 악성코드다/정상파일이다 정답을 표시해 놓는 것이다. 모델은 이를 통해 패턴을 인식하고 학습할 수 있다.

라벨링은 바이러스 토탈의api를 사용하여 진행한다.

바이러스 토탈에 접속하여 아이디가 없다면 회원가입을 한다.

https://www.virustotal.com/gui/home/upload

회원가입 후 프로필 사진을 클릭해 API Key로 들어간다.

바이러스 토탈 api key를 저장해 둔다.

1.1. 라벨링 자동화 코드 작성

파이썬 코드를 작성하여 많은 수의 데이터 라벨링 작업을 자동화하도록 한다.

실습자료에서 제공하는 코드가 있지만 python 3, anaconda3, virustotl api 버전과 호환되지 않아 새로 코드를 작성하였다.

import json
import hashlib
import os
import re
import time
import requests

# Regular expression to check if a filename matches an MD5 hash pattern
md5_pattern = re.compile(r'[a-f0-9]{32}')
path_dir = '/home/stud/sample'  # Directory with sample files
file_list = os.listdir(path_dir)  # List of files in the directory

class VirusTotalAPI:
    def __init__(self):
        self.api_key = '07319ec595124ff48fe0032e4fb70144afa8937b90d96716985f7b49e1ece07a'  # Hardcoded API Key
        self.base_url = 'https://www.virustotal.com/api/v3/'

    def get_report(self, md5):
        url = f"{self.base_url}files/{md5}"
        headers = {"x-apikey": self.api_key}
        response = requests.get(url, headers=headers)

        if response.status_code == 200:
            jdata = response.json()
            positives = jdata['data']['attributes']['last_analysis_stats']['malicious']
            print(f"=== Results for MD5: {md5}\tDetected by: {positives}")
            return positives
        elif response.status_code == 404:
            print(f"{md5} -- Not Found in VirusTotal, requesting a scan")
            return None
        else:
            print(f"Error: {response.status_code} - {response.json().get('error', {}).get('message', 'Unknown error')}")
            return None

    def request_scan(self, filepath):
        url = f"{self.base_url}files"
        headers = {"x-apikey": self.api_key}
        files = {"file": (os.path.basename(filepath), open(filepath, "rb"))}

        response = requests.post(url, headers=headers, files=files)
        if response.status_code == 200:
            jdata = response.json()
            return jdata['data']['id']  # Returns file ID for scan tracking
        else:
            print(f"Error during scan request: {response.status_code} - {response.json().get('error', {}).get('message', 'Unknown error')}")
            return None

    @staticmethod
    def get_md5(filepath, blocksize=8192):
        md5 = hashlib.md5()
        with open(filepath, "rb") as f:
            while chunk := f.read(blocksize):
                md5.update(chunk)
        return md5.hexdigest()

def main():
    vt = VirusTotalAPI()
    for file in file_list:
        filepath = os.path.join(path_dir, file)
        # Check if filename is an MD5 hash
        if not md5_pattern.fullmatch(file):
            file = vt.get_md5(filepath)  # Calculate MD5 hash if filename is not MD5 format
        # Attempt to get report or request a scan
        positives = vt.get_report(file)
        if positives is None:
            scan_id = vt.request_scan(filepath)
            if scan_id:
                print(f"Waiting for scan results for {file}...")
                time.sleep(15)  # Wait time to check scan status
                while True:
                    positives = vt.get_report(file)
                    if positives is not None:
                        break

        # Rename file according to detection result
        new_filename = f"{positives}#{file}"
        new_filepath = os.path.join(path_dir, new_filename)
        os.rename(filepath, new_filepath)
        print(f"Renamed {filepath} to {new_filepath}")

if __name__ == '__main__':
    main()

코드 개요

/home/stud/sample에 저장된 데이터셋을 대상으로 진행된다.
클래스에는 get_report(), request_scan(), get_md5() 메서드가 포함되어 있으며, 각각 파일의 악성 여부를 확인하고, 필요한 경우 스캔을 요청하며, 파일의 MD5 해시를 계산한다.
라벨링 된 파일의 이름은 [탐지 엔진 수(positives) + # + MD5 해시 형식] 이다.
예를 들어 스캔 결과로 탐지 엔진 수가 5, MD5 해시가 abcd1234...라고 하면, 파일 이름은 5#abcd1234... 형식이 된다. 정상프로그램일 경우 "0#md5해시값" 이 된다
클래스에는 get_report(), request_scan(), get_md5() 메서드가 포함되어 있으며, 각각 파일의 악성 여부를 확인하고, 필요한 경우 스캔을 요청하며, 파일의 MD5 해시를 계산합니다.
virustotal에서 api 호출을 1분에 4개로 제한하기 때문에 분석된적없는 새로운 파일을 라벨링 할경우 15초 간격으로 메소드를 호출하도록 한다.

메서드의 기능

파일 악성 검사 (get_report 메서드)

get_report()는 파일의 MD5 해시를 기반으로 VirusTotal에서 파일의 스캔 결과를 가져온다.
VirusTotal에서 파일이 이미 분석된 적이 있다면, 반환된 JSON 데이터에서 악성으로 탐지된 엔진 수(positives)를 확인한다.
만약 해당 MD5 해시가 데이터베이스에 존재하지 않으면 None을 반환하여 스캔 요청을 위한 준비를 한다.

새로운 스캔 요청 (request_scan 메서드)

request_scan()는 데이터베이스에 없는 파일을 VirusTotal에 업로드하여 새로운 스캔을 요청한다.
요청이 성공적으로 수행되면 스캔 추적을 위한 파일 ID가 반환된다.

파일 MD5 해시 계산 (get_md5 메서드)

get_md5() 메서드는 파일을 읽어 MD5 해시를 계산하여 반환한다.
파일 이름이 이미 MD5 해시 형식이 아닌 경우 이 메서드를 호출하여 고유한 MD5 해시 값을 생성한다.

1.2. 자동화 코드 테스트

방금 전 발급받았던 api키로 파이썬 코드를 수정하고 "virustotal_labeling.py" 라는 이름으로 저장한다.

그리고 실습자료 중 샘플코드를 준비해 놓는다. 경로는 다음과 같다.

https://github.com/bjpublic/Ai-learnsecurity

다음으로 우분투 가상머신을 열어 home 폴더에 sample , scripts 이름을 가진 폴더를 각각 생성한다.

파이썬 코드는 우분투 scripts 폴더에 복사, 실습 샘플코드는 sample 폴더에 복사한다.

home 폴더에서 터미널을 열어 아래 명령어를 입력한다.

conda activate mlsec_3811

pip install requests

scripts 폴더로 이동 파이썬 코드 실행

cd /scripts

python3 /home/stud/scripts/virustotal_labeling.py

결과가 잘 나왔다.

1.3. 데이터셋 라벨링

이제 사용할 악성코드/정상프로그램 데이터셋을 라벨링 한다.

무료 api사용은 1분에 4개 질의 제한에 하루 제한 500개가 있으므로 1000개의 데이터를 라벨링 하기위 해서는 꼬박 이틀이 걸린다.

참고

책에서 제공하는 샘플파일중 정상프로그램에 중복된 파일이있어 라벨링 할경우 최종 파일 갯수가 몇개 줄어든다.

아래 파일들이 중복된 파일이다.

[인공지능 보안을 배우다] 7장 PJ_1 악성코드 탐지모델 (지도학습을 위한 레이블링)

1. 레이블링 이란

1.1. 라벨링 자동화 코드 작성

1.2. 자동화 코드 테스트

1.3. 데이터셋 라벨링

참고

티스토리툴바