[Python] VirusTotal API + Response Save Bot 만들기

Python을 활용한 업무의 일부를 공유하는 게시물입니다.

아직 공부하는 입장이기에 부족한 부분이 있을 수 있습니다. 지적해주시면 적극 반영하도록 하겠습니다.
퍼가시거나, 업무에서 도움이 되었다면 댓글 하나씩 부탁드립니다 :) 감사합니다.

---

Summary

어떤 파일이 있을 때, 해당 파일의 악성/정상 유무를 확인하기 가장 손 쉬운 방법은 VirusTotal에 MD5 기준으로 질의하는 것이다.

하지만, 확인해야되는 파일 개수가 너무 많으면 사람이 처리할 수 없고, 이때 VirusTotal API를 사용하면 아주아주 편하다.

VirusTotal

EXE 파일로 변환하면 아래와 같은 과정으로 File Hash를 읽고, 이를 통해 악성코드 유/무를 확인한다.

1. 특정 경로에 위치한 txt 파일을 불러오기

초기 화면

2. 해당 txt 파일에서 문자열 길이 기준으로 파일 MD5 식별

3. 행을 1개의 파일 MD5로 인지하여, Virutotal API 질의

TXT 파일 읽기

4. API Response에서 유명 악성코드 벤더의 결과를 엑셀로 자동 저장(분석이 완료도면 Popup 발생)
- 무료 API Key의 경우 1분에 3개까지만 MD5 조회할 수 있으므로 각 API 요청마다 20s Sleep 추가

결과 팝업

결과 화면

Python Code

혹시나 필요하신 분이 있으시면 아래 Code에서 VT_KEY 변수만 수정해서 사용하시면 됩니다~

 - Code도 어렵지 않으므로, 확인할 악성코드 벤더를 조정하거나 하시면 됩니다~

#! /usr/bin/env python
# -*- coding: UTF-8 -*-
# coding=utf-8

import ctypes
import os
import ssl
import threading
import time
import urllib
from datetime import datetime
from tkinter import *
from tkinter import filedialog
from tkinter import ttk as ttk
import openpyxl
import requests

#create user application
filename = 'name'
aa = 'tmp'
window=Tk()
window.title('비인가 프로그램 검증 봇')
window.geometry("500x70")
main=Menu(window)
window.config(menu=main)
fileMenu=Menu(main,tearoff=0) #tearoff=0을 해줘야 기본 ---- 줄이 없어짐
helper=Menu(main,tearoff=0)

class guidemo:
    def end():
        window.quit()
        window.destroy()

    def new():
        global filename
        filename=filedialog.askopenfilename(parent=window, filetype=(('txt파일','*.txt'), ('모든파일','*.*')))

        if filename[-3:] == 'txt':
            tmp = threading.Thread(target = vt_req, daemon = True) # data thread create.. maybe?
            # daemon 설정은 하니까 main 죽으면 Thread도 같이 죽는다
            # daemon은 수명이 메인 스레드에 의존적이므로, 메인 스레드가 종료될 경우 자신도 종료되고 리소스를 반납한다.
            tmp.start()
        else:
            msg = ctypes.windll.user32.MessageBoxW(None, ".txt 파일을 다시 선택해주세요", "오류", 0)

    def help():
        msg = ctypes.windll.user32.MessageBoxW(None, "본 프로그램은 Virustotal API을 이용한 분석 자동화 프로그램입니다.\n.txt 파일 내 1줄을 MD5로 인지하고 API 요청을 하게 됩니다.\n\n사용 예시)\n(1)C:/TEST/TEST.txt 파일을 읽음\n(2)txt 파일 내 MD5를 Virustotal API를 사용하여 분석 요청\n(3)분석 결과를 동일한 경로에 xlsx형태로 저장\n\n\n!!!해당 프로그램 사용으로 발생하는 문제는 책임지지 않습니다.!!!", "사용설명서", 0)

    def __init__(self):
        #window.iconbitmap(default = r'C:\Users\Dongjin\Pictures\img.ico') -> 이걸하면 exe 만들고 다른 PC에서 해당 리소스가 없어서 에러 남
        main.add_cascade(label='파일', menu=fileMenu)
        main.add_cascade(label='도움말', menu=helper)
        fileMenu.add_command(label='열기', command=guidemo.new)
        fileMenu.add_separator()
        fileMenu.add_command(label='종료', command=guidemo.end)
        helper.add_command(label='도움말', command=guidemo.help)
        window.mainloop()


def search_result(scan, key, vendor):
    global aa
    aa="tmp"
    if key == vendor:
        aa = str(scan[key]['result'])
        if aa == 'None':
            aa = "Normal"
    return aa


def vt_req():
    VT_KEY = '바이러스토탈 API 키 정보' <<< 이 부분을 수정하시면 됩니다.
    REPORT_URL = "https://www.virustotal.com/vtapi/v2/file/report"
    context = ssl._create_unverified_context()
    row_v = 2
    cnt = 0 # 분석 완료 후 알람을 띄어주기 위한 변수

    #read file hash
    fd = open(filename,'r')
    fd_2 = open(filename,'r') # to search line in txt file
    num = list(enumerate(fd_2))[-1][0]+1 #calculate line in txt file
    des = f'## {filename}에서 확인 된 분석 요청 개수는 총 {num}개 입니다. ##'
    an_time = num * 20
    des_2 = f'## 소요 시간은 {an_time}초 예상됩니다. ##'

    #show application data
    label = ttk.Label(window, text = des)
    label2 = ttk.Label(window, text = des_2)
    label.pack()
    label2.pack()

    #creat output file
    s = datetime.today().strftime("%Y%m%d")
    wb = openpyxl.Workbook()
    ws = wb.active
    ws.cell(row=1, column=1).value = 'MD5'
    ws.cell(row=1, column=2).value = 'AhnLab-V3'
    ws.cell(row=1, column=3).value = 'FireEye'
    ws.cell(row=1, column=4).value = 'Kaspersky'
    ws.cell(row=1, column=5).value = 'Sophos'
    ws.cell(row=1, column=6).value = 'TrendMicro'
    ws.cell(row=1, column=7).value = 'Microsoft'
    ws.cell(row=1, column=8).value = 'McAfee'
    ws.cell(row=1, column=9).value = 'Fortinet'
    ws.cell(row=1, column=10).value = 'Avast'
    ws.cell(row=1, column=11).value = 'BitDefender'

#    print ("----------------------------------------")
#    print ('!!!!!Virus Total API Scanning Start!!!!!')

    while True :
        line = fd.readline()
        md5str = line.strip('\n')
        if not md5str : break
#        print (md5str)
#        print (len(md5str))
        if len(md5str) != 32 :
            ws.cell(row=row_v, column=1).value = md5str
            ws.cell(row=row_v, column=2).value = 'File Hash is invalid.'
            row_v = row_v + 1
            cnt += 1
            if cnt == num:
                des_3 = f'######## 분석이 완료 되었습니다. #######'
                label3 = ttk.Label(window, text = des_3)
                label3.pack()
                msg = ctypes.windll.user32.MessageBoxW(None, "모든 VirusTotal 분석 요청이 완료 되었습니다.", "완료", 0)
        else:
            #scanning start
            para = {'resource' : md5str, 'apikey' : VT_KEY}
            en_para = urllib.parse.urlencode(para)
            time.sleep(2)
            os.environ['http_proxy']='' #slove proxy problem
            res = requests.get(REPORT_URL, params=en_para)
            data = res.json()
            md5 = data.get('md5', {})
            scan = data.get('scans', {})
            keys = scan.keys()
#            print ("  ")
#            print ("----------------scanning---------------")
#            print (md5str)
            ws.cell(row=row_v, column=1).value = md5str
#            print ("---------------------------------------")
            if md5 == {}:
#                print ("!!!!! sorry no match !!!!!")
                ws.cell(row=row_v, column=2).value = 'There is no file in the VT.'

            # 인지도 있는 10개의 벤더 결과 출력
            for key in keys :
                result = search_result(scan, key, 'AhnLab-V3')
                if aa !="tmp" :
                    ws.cell(row=row_v, column=2).value = result
                result = search_result(scan, key, 'FireEye')
                if aa !="tmp" :
                    ws.cell(row=row_v, column=3).value = result
                result = search_result(scan, key, 'Kaspersky')
                if aa !="tmp" :
                    ws.cell(row=row_v, column=4).value = result
                result = search_result(scan, key, 'Sophos')
                if aa !="tmp" :
                    ws.cell(row=row_v, column=5).value = result
                result = search_result(scan, key, 'TrendMicro')
                if aa !="tmp" :
                    ws.cell(row=row_v, column=6).value = result
                result = search_result(scan, key, 'Microsoft')
                if aa !="tmp" :
                    ws.cell(row=row_v, column=7).value = result
                result = search_result(scan, key, 'McAfee')
                if aa !="tmp" :
                    ws.cell(row=row_v, column=8).value = result
                result = search_result(scan, key, 'Fortinet')
                if aa !="tmp" :
                    ws.cell(row=row_v, column=9).value = result
                result = search_result(scan, key, 'Avast')
                if aa !="tmp" :
                    ws.cell(row=row_v, column=10).value = result
                result = search_result(scan, key, 'BitDefender')
                if aa !="tmp" :
                    ws.cell(row=row_v, column=11).value = result
            row_v = row_v + 1
            time.sleep(20) #sleep time because VT API accept 4 file each 1 min
            cnt += 1
            if cnt == num:
                des_3 = f'######## 분석이 완료 되었습니다. ########'
                label3 = ttk.Label(window, text = des_3)
                label3.pack()
                msg = ctypes.windll.user32.MessageBoxW(None, "모든 VirusTotal 분석 요청이 완료 되었습니다.", "완료", 0)

    #file save and close
    wb.save(s + '_파일_분석_결과.xlsx')
    wb.close()
    fd.close()

if __name__ == "__main__":
    guidemo()

끝!!