import os
import requests
import time

requests.adapters.DEFAULT_RETRIES = 5
url = 'https://www.virustotal.com/vtapi/v2/file/report'
params = {'apikey': 'b1571259fa877a755620a3cacb5833622fcdd8695a81be93306a747c408795e7',
          'resource': '<resource>'}

output = {}

if os.path.exists('result_pku.json'):
    fopen = open('result_pku.json','r')
    output = eval(fopen.read())
    fopen.close()


md52file = {}

fopen = open('dic_pku.json','r')
md52file = eval(fopen.read())
fopen.close()

for idx,md5 in enumerate(md52file):
    if md5 in output:
        # print('md5', md5)
        print('pass')
        continue
    print('md5',md5)
    dir = md52file[md5]['dir']
    dir = dir.split('\\')[1]
    print('raw_label',dir)
    # if not dir == 'Trojan':
    #     continue
    params['resource'] = md5

    response = requests.get(url, params=params)
    while response.status_code != 200:
        print('未获取信息')

        time.sleep(3)
        response = requests.get(url, params=params)

    report = response.json()
    print(report)
    if 'scans' in report:
        if 'Kaspersky' in report['scans']:
            print(report['scans']['Kaspersky']['result'])
            label = report['scans']['Kaspersky']['result']
            if not label is None:
                label = label.split('.')[0]
                label = label.split(':')[-1]
            print(md5, label)
            output[md5] = {}
            output[md5]['label'] = label
            output[md5]['dir'] = md52file[md5]['dir']

    if idx % 10 == 0:
        fwrite = open('result_pku.json', 'w')
        fwrite.write(str(output))
        fwrite.close()

print(output)