[Virustotal API v3] MITRE ATT&CK 파싱 코드

이전 프로젝트에 사용하기 위해 만들었던 스크립트입니다.

---

import requests
import hashlib
import time
import json

# 예시 정답
answer = "T1027"

# 사용자가 제출한 코드에서 추출한 techniques 번호들을 모을 리스트
techniques = []

# VirusTotal API 키
VT_API_KEY = "" #개인 VT 키

# VirusTotal API 엔드포인트
VT_API_SCAN_URL = "https://www.virustotal.com/vtapi/v2/file/scan"
VT_API_URL = "https://www.virustotal.com/api/v3/files"

def calculate_md5(file_path):
    # 파일 내용을 읽어와 MD5 해시 계산
    with open(file_path, 'rb') as file:
        md5_hash = hashlib.md5(file.read()).hexdigest()
    return md5_hash

def upload_to_virustotal(file_path):
    # 파일의 MD5 해시 계산
    md5_hash = calculate_md5(file_path)

    # 파일을 업로드하기 위한 POST 요청
    files = {'file': (file_path, open(file_path, 'rb'))}
    params = {'apikey': VT_API_KEY}
    
    response = requests.post(VT_API_SCAN_URL, files=files, params=params)

    # 응답을 JSON 형식으로 변환
    try:
        result = response.json()

        # 업로드 결과 출력
        if result['response_code'] == 1:
            print("File successfully uploaded to VirusTotal.")
            print(f"MD5 Hash: {md5_hash}")
            print("VirusTotal Scan ID:", result['scan_id'])

            # 보고서 얻기
            get_report(md5_hash)
        else:
            print("File upload to VirusTotal failed.")
            print("Response Code:", result['response_code'])
    except ValueError as e:
        print(f"Error decoding JSON response from VirusTotal: {e}")

def get_report(md5_hash):
    # 보고서를 얻기 위한 GET 요청
    VT_API_REPORT_URL = VT_API_URL + "/" + str(md5_hash) + "/behaviour_mitre_trees"

    headers = {"accept": "application/json","x-apikey": VT_API_KEY}
    response = requests.get(VT_API_REPORT_URL, headers=headers)

    # 응답을 JSON 형식으로 변환
    try:
        report = response.json()

        # JSON 파일로 저장
        save_to_json(md5_hash, report)

    except ValueError as e:
        print(f"Error decoding JSON response from VirusTotal: {e}")

def save_to_json(md5_hash, report):
    # 파일 이름은 MD5 해시를 사용
    file_name = f"{md5_hash}.json"
    json_file_path = "C:/Users/Owner/Desktop/"+file_name

    with open(json_file_path, 'w') as json_file:
        json.dump(report, json_file, indent=4)

    print(f"VirusTotal report saved to {file_name}")
    parshing_json(md5_hash, json_file_path)

def parshing_json(md5_hash, json_file_path):
    file_name = f"{md5_hash}.txt"
    output_file_path = "C:/Users/Owner/Desktop/"+file_name

    with open(json_file_path, 'r') as file:
        data = json.load(file)

    # virustotal report에서 tecniques 번호를 파싱해서 바탕화면에 txt 파일로 저장
    with open(output_file_path, 'w') as output_file:
        for tool_name, tool_data in data['data'].items():
            for tactic in tool_data['tactics']:
                for technique in tactic['techniques']:
                    technique_id = technique['id'] 
                    output_file.write(f"{technique_id}\n") 
                    techniques.append(technique_id) # 정답 비교를 위해 리스트에 technique 번호 저장
    print(f"VirusTotal report techniques saved to: {output_file_path}")
    if answer in technique_id: # techniques 리스트에 정답이 있다면 정답, 불일치하면 오답
        print("정답입니다")
    else:
        print("오답입니다")

def main():
    #업로드할 코드(실행파일)의 경로
    file_path = ""

    # VirusTotal에 파일 업로드
    upload_to_virustotal(file_path)

if __name__ == "__main__":
    main()

 

API v2는 더 이상 사용되지 않으며, API v3이 사용됨 => API v3로는 파일 업로드가 안되며, 문제를 해결하지 못해 v2로 업로드