malware.py

from .config import *
from .utils import format_dict
import os
import time
import requests


class VirusTotal:

    def __init__(self):
        self.params = {'apikey': virus_total_api_key}
        self.scan_url = 'https://www.virustotal.com/vtapi/v2/file/scan'
        self.report_url = 'https://www.virustotal.com/vtapi/v2/file/report'

    @staticmethod
    def check_response(response):
        if response.status_code == 204:
            print('Request rate limit exceeded. Try after some time.')
        elif response.status_code == 403:
            print("Forbidden. You don't have enough privileges to make the request.\nCheck your API Key.")
        else:
            return True
        return False

    def send_malware(self, file_name, json_output=False):
        try:
            files = {'file': (file_name, open(file_name, 'rb'))}
        except FileNotFoundError:
            print(f'File {file_name} not found.')
        else:
            max_size = 32*(1024**2)  # 32 MB
            if os.stat(file_name).st_size > max_size:
                print(f'File Size can not be greater than 32 MB.')
            response = requests.post(self.scan_url, files=files, params=self.params)
            if VirusTotal.check_response(response):
                json_response = response.json()

                if json_output:  # --json format
                    print(json_response)
                else:
                    format_dict(json_response)
                print("\nMalware sample is submitted!! Wait for 5 minutes !!"
                      " Please do not stop this script !!\n")
                time.sleep(300)

                # Getting the report back
                resource = json_response['resource']
                self.params['resource'] = resource
                headers = {"Accept-Encoding": "gzip,deflate",
                           "User-Agent": "gzip, OSINT-SPY"}
                response = requests.get('https://www.virustotal.com/vtapi/v2/file/report',
                                        params=self.params, headers=headers)
                if VirusTotal.check_response(response):
                    json_response = response.json()

                    if json_output:  # --json
                        print(json_response)
                    else:
                        format_dict(json_response)