# !/bin/python
# -*- coding: utf-8 -*-
u"""SecureTea Social Engineering
Project:
    ╔═╗┌─┐┌─┐┬ ┬┬─┐┌─┐╔╦╗┌─┐┌─┐
    ╚═╗├┤ │  │ │├┬┘├┤  ║ ├┤ ├─┤
    ╚═╝└─┘└─┘└─┘┴└─└─┘ ╩ └─┘┴ ┴
    Author: Digvijay Bhosale <digvijayb1729@gmail.com> August 16 2021
    Version: 1.0
    Module: SecureTea
"""

import subprocess
import os
import json
import time
import sys
import threading
from securetea.lib.malware_analysis import malwareAnalysisJSONDisplay
from securetea.lib.malware_analysis import globals

'''
Note : VirusTotal API limit is 4 req per min and 500 req per day.
'''


class MalwareAnalysis:
    '''
    Contains functions that upload files and analyse them using VirusTotal public (free) API

    runner()
    This function is responsible for running other functions and as such, program flow starts here

    virustotal_file_upload()
    Uploads a file to virustotal

    virustotal_get_results()
    Downloads results of analysis

    show_results()
    Displays result according to user verbosity
    '''

    def __init__(self, filename=None, API_KEY=None):
        self.filename = filename
        self.API_KEY = API_KEY
        self.report_id = ''
        self.results = ''
        self.av_engines_list = ''

        globals.initialize_colours()

    def __del__(self):
        self.report_id = ''
        self.results = ''
        self.av_engines_list = ''

    def runner(self):
        self.virustotal_file_upload()
        self.virustotal_get_results()
        self.show_results()
        change_verbosity = 'y'
        while change_verbosity.lower() != 'n':
            change_verbosity = input('Would you like to see results with changed verbosity ? (y/N) \n\t: ') or 'n'
            if change_verbosity.lower() == 'y':
                self.show_results()
            else:
                break

    def threat_level(self):
        self.virustotal_file_upload()
        self.virustotal_get_results(noisy=False)
        display_obj = malwareAnalysisJSONDisplay.JSONDisplay()
        detected, undetected, unsupported, total = display_obj.threat_level()
        return detected, undetected, unsupported, total

    def threat_level_threading(self):
        """
        this is called in case of continuous file analysis threading version
        uploads files, downloads reports
        Sends to JSON program
        """
        self.virustotal_file_upload()
        self.virustotal_get_results(noisy=False)
        display_obj = malwareAnalysisJSONDisplay.JSONDisplay(self.av_engines_list)
        display_obj.threat_level_threading(self.filename, self.report_id)
        # the threat_level_threading function ^ is from display program. Its not recursing

    def virustotal_file_upload(self, filename=None, API_KEY=None):
        if filename is not None:
            self.filename = filename
        if API_KEY is not None:
            self.API_KEY = API_KEY
        file_size = os.path.getsize(self.filename)
        if int(file_size) < 32000000:
            base_url = 'https://www.virustotal.com/api/v3/files'
            headers = 'x-apikey: ' + self.API_KEY
            file_path = '@' + self.filename

            process = subprocess.Popen(['curl', '--request', 'POST', '--url', base_url,
                                        '--header', headers,
                                        '--form', 'file=' + file_path],
                                       stdout=subprocess.PIPE,
                                       stderr=subprocess.PIPE)
            stdout, stderr = process.communicate()
            data = stdout.decode('utf-8')
            stderr = stderr.decode('utf-8')
            json_data = json.loads(data)
            self.report_id = json_data['data']['id']

        else:
            print(globals.OKBLUE + 'File Size greater than 32MB' + globals.END)
            print(globals.OKBLUE + 'Current file size is :' + str(file_size) + globals.END)
            print(globals.OKBLUE + 'Kindly visit https://www.virustotal.com/ and submit file' + globals.END)

    def virustotal_get_results(self, noisy=True, report_id=None, API_KEY=None):
        if report_id is not None:
            self.report_id = report_id
        if API_KEY is not None:
            self.API_KEY = API_KEY
        if self.report_id == '':
            print('No report ID available')
            return
        base_url = 'https://www.virustotal.com/api/v3/analyses/'
        url = base_url + str(self.report_id)
        headers = 'x-apikey: ' + self.API_KEY
        if noisy:
            print(globals.OKCYAN + 'Waiting for results from VirusTotal. This might take some time. (Typically 0.5-4 Minutes)' + globals.END)
        self.results = self.curl_results(url, headers, noisy)
        self.av_engines_list = self.results['data']['attributes']['results']

    def curl_results(self, url, headers, noisy=True):
        self.waiting_and_animation()
        process = subprocess.Popen(['curl', '--request', 'GET', '--url', url,
                                    '--header', headers],
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE)
        stdout, stderr = process.communicate()
        data = stdout.decode('utf-8')
        stderr = stderr.decode('utf-8')
        res = json.loads(data)
        if res['data']['attributes']['status'] == 'queued':  # using recursive funcs to get complete results
            if noisy:
                print('Status : ' + str(res['data']['attributes']['status']))
            res = self.curl_results(url, headers, noisy=noisy)            # if its queued, wait(30)
            return res
        else:
            print(globals.HEADER + 'Status : ' + str(res['data']['attributes']['status']) + globals.END)
            return res

    def waiting_and_animation(self):
        animation_time = 0
        while animation_time < 30:
            anim = "\|/-\|/-"
            for char in anim:
                sys.stdout.write(char)
                sys.stdout.flush()
                sys.stdout.write('\b')
                time.sleep(0.1)
                animation_time = animation_time + 0.1

    def show_results(self):
        if self.results == '':
            print(globals.FAIL + 'No results available' + globals.END)
            return
        try:
            verbosity = int(input('Enter Verbosity (Default : 2) : ') or 2)
        except KeyboardInterrupt:
            print(globals.OKCYAN + '\nKeyboardInterrupt..... Quitting' + globals.END)
            return
        except TypeError:
            verbosity = 2

        display_obj = malwareAnalysisJSONDisplay.JSONDisplay(self.av_engines_list)
        if verbosity == 1:
            display_obj.verbosity_1()
        elif verbosity == 2:
            display_obj.verbosity_2()
        elif verbosity == 3:
            display_obj.verbosity_3()
        elif verbosity == 4:
            display_obj.verbosity_4()
        else:
            print(globals.OKCYAN + "Given Verbosity not available" + globals.END)
            display_obj.verbosity_1()
