VirusTotal.py

import logging

import requests

from django.conf import settings

from mobsf.MobSF.proxy import upstream_proxy

from mobsf.MobSF.utils import (
    append_scan_status,
    file_size,
    get_config_loc,
)

logger = logging.getLogger(__name__)


class VirusTotal:

    API_KEY_ERROR = 'VirusTotal Permission denied, wrong api key?'
    API_KEY_ERROR_SHORT = 'VirusTotal API error'
    API_CONN_ERROR = 'VirusTotal Connection Error'

    def __init__(self, checksum):
        self.base_url = settings.VIRUS_TOTAL_BASE_URL
        self.checksum = checksum

    def get_report(self):
        """
        Get Report from VT.

        :param file_hash: md5/sha1/sha256
        :return: json response / None
        """
        try:
            file_hash = self.checksum
            url = self.base_url + 'report'
            params = {
                'apikey': settings.VT_API_KEY,
                'resource': file_hash}
            headers = {'Accept-Encoding': 'gzip, deflate'}
            try:
                proxies, verify = upstream_proxy('https')
            except Exception as exp:
                msg = 'Setting upstream proxy'
                logger.exception(msg)
                append_scan_status(file_hash, msg, repr(exp))
            try:
                response = requests.get(
                    url,
                    timeout=5,
                    params=params,
                    headers=headers,
                    proxies=proxies,
                    verify=verify)
                if response.status_code == 403:
                    logger.error(self.API_KEY_ERROR)
                    append_scan_status(
                        file_hash,
                        self.API_KEY_ERROR,
                        self.API_KEY_ERROR_SHORT)
                    return None
            except Exception as exp:
                logger.error(self.API_CONN_ERROR)
                append_scan_status(
                    file_hash,
                    self.API_CONN_ERROR,
                    repr(exp))
                return None
            try:
                json_response = response.json()
                return json_response
            except ValueError:
                return None
        except Exception as exp:
            msg = 'Failed to get report from VirusTotal'
            logger.exception(msg)
            append_scan_status(file_hash, msg, repr(exp))
            return None

    def upload_file(self, file_path):
        """
        Upload File to VT.

        :param file_path: file path to upload
        :return: json response / None
        """
        try:
            url = self.base_url + 'scan'
            if file_size(file_path) > 31:
                logger.warning('VirusTotal Public API does '
                               'not support files above 32 MB')
                return None
            files = {'file': open(file_path, 'rb')}
            headers = {'apikey': settings.VT_API_KEY}
            try:
                proxies, verify = upstream_proxy('https')
            except Exception as exp:
                msg = 'Setting upstream proxy'
                logger.exception(msg)
                append_scan_status(self.checksum, msg, repr(exp))
            try:
                response = requests.post(
                    url,
                    timeout=5,
                    files=files,
                    data=headers,
                    proxies=proxies,
                    verify=verify)
                if response.status_code == 403:
                    logger.error(self.API_KEY_ERROR)
                    append_scan_status(
                        self.checksum,
                        self.API_KEY_ERROR,
                        self.API_KEY_ERROR_SHORT)
                    return None
            except Exception as exp:
                logger.error(self.API_CONN_ERROR)
                append_scan_status(
                    self.checksum,
                    self.API_CONN_ERROR,
                    repr(exp))
                return None
            json_response = response.json()
            return json_response

        except Exception as exp:
            msg = 'Failed to upload file to VirusTotal'
            logger.exception(msg)
            append_scan_status(self.checksum, msg, repr(exp))
            return None

    def get_result(self, file_path):
        """
        Get Results from VT.

        Uploading a file and getting the approval msg from VT
        or fetching existing report
        :param file_path: file's path
        :return: VirusTotal result json / None upon error
        """
        try:
            file_hash = self.checksum
            msg = 'VirusTotal: Check for existing report'
            logger.info(msg)
            append_scan_status(file_hash, msg)
            report = self.get_report()
            # Check for existing report
            if report:
                if report['response_code'] == 1:
                    msg = f'VirusTotal: {report["verbose_msg"]}'
                    logger.info(msg)
                    append_scan_status(file_hash, msg)
                    return report
            if settings.VT_UPLOAD:
                msg = 'VirusTotal: file upload'
                logger.info(msg)
                append_scan_status(file_hash, msg)
                upload_response = self.upload_file(file_path)
                if upload_response:
                    msg = f'VirusTotal: {upload_response["verbose_msg"]}'
                    logger.info(msg)
                    append_scan_status(file_hash, msg)
                return upload_response
            else:
                msg = ('VirusTotal Scan not performed as file '
                       f'upload is disabled in {get_config_loc()}. '
                       'To enable file upload, set VT_UPLOAD to True.')
                logger.info(msg)
                append_scan_status(file_hash, msg)
                message = ('Scan not performed, VirusTotal file '
                           f'upload disabled in {get_config_loc()}')
                report = {
                    'verbose_msg': message,
                    'positives': 0,
                    'total': 0}
                return report
        except Exception as exp:
            msg = 'VirusTotal: Error getting result'
            logger.exception(msg)
            append_scan_status(file_hash, msg, repr(exp))