import requests
import json
import pyzipper
import os
import sys


url = 'https://mb-api.abuse.ch/api/v1/'
password = b'infected'


def get_path(value):
    if value[len(value)-1] != '/':
        value = value + '/'
    return value


def get_list(query, **kwargs):
    data = {'query': query}
    for kw in kwargs:
        data[kw] = kwargs[kw]
    print(data)
    try:
        response = requests.post('https://mb-api.abuse.ch/api/v1/', data=data)
    except Exception as e:
        print(e)
    return response.text


def get_list_old(type):
    if type == 'limit':
        type = '100'
    else:
        type = 'time'
    data = {'query': 'get_recent', 'selector': type}
    try:
        response = requests.post('https://mb-api.abuse.ch/api/v1/', data=data)
    except Exception as e:
        print(e)
    return response.text


def make_hash_list(text):
    malware_json = json.loads(text)
    list_of_hashes = []
    for hash in malware_json['data']:
        list_of_hashes.append(hash['sha256_hash'])
    print(list_of_hashes)
    return list_of_hashes


def retrieve_sample_to_disk(url, path, hash):
    data = {'query': 'get_file', 'sha256_hash': str(hash)}
    try:
        response = requests.post(url, data=data)
    except Exception as e:
        print(e)
    with open(path + hash, 'wb') as f:
        f.write(response.content)


def write_cleartext_files(path, list_of_shas):
    for item in list_of_shas:
        with pyzipper.AESZipFile(path + item) as zf:
            zf.setpassword(password)
            zf.extractall(path + 'unpacked/')
            print(f'Finished unpacking into {item}')
        file_path = path + item
        try:
            os.remove(file_path)
            print(f'Removed zip file {item}')
        except Exception as e:
            print(e)


def main():
    # sys.argv[1] is the path for output, we need this
    if len(sys.argv) > 1:
        path = get_path(sys.argv[1])
    else:
        print('''First argument must be path for malware output, 
                second number of each type of sample''')
        exit(1)
    # If there is a second argument use it as the number of samples
    if len(sys.argv) > 2:
        number = str(sys.argv[2])
    else:
        number = '1'
    # Go get 'em. Tags to download are hardcoded to avoid infinte args
    for tag in ['exe', 'dll', 'docx', 'xlsx', 'jar', 'elf']:
        tag_path = path + tag + '/'
        os.makedirs(tag_path, exist_ok=True)
        # This just gets the result from Malware bazaar for the
        # type of query you pass it
        malware_list = get_list('get_taginfo', tag=tag, limit=number)
        # this turns it into a list of hashes to request
        list_for_downloads = make_hash_list(malware_list)
        # this gets them and writes them to disk
        for hash in list_for_downloads:
            retrieve_sample_to_disk(url, tag_path, hash)
        tag_path_unpacked = tag_path + 'unpacked/'
        os.makedirs(tag_path_unpacked, exist_ok=True)
        write_cleartext_files(tag_path, list_for_downloads)


if __name__ == '__main__':
    main()
