malware_detect.py

import os
import json
import datetime
import base64
import uuid
import shutil
from mimetypes import MimeTypes

import requests
from flask import make_response, request
from flask_restful import Resource
from clustering import clustering

UPLOAD_DIR = 'uploads'


class CustomException(Exception):
    pass


class StillProcessing(Exception):
    pass


class MalwareDetect(Resource):

    def __init__(self):
        self.session = requests.Session()
        self.headers = {
            'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) '
                          'Chrome/74.0.3729.169 Safari/537.36',
        }
        self.user_data = None
        self.device_id = uuid.uuid4().hex
        self.uuid = uuid.uuid4().hex
        self.mime = MimeTypes()
        current_time = int(datetime.datetime.utcnow().timestamp())
        self.upload_filename = f'{current_time}.json'
        print(self.upload_filename)
        self.response = {'malware': '', 'confidence': 0.0}

    def post(self):
        response_object = {
            'status': 'fail',
            'message': 'Invalid payload.'
        }
        is_parse = request.is_json
        if not is_parse:
            response_object['message'] = 'Arguments parsing is failed.'
            return make_response(json.dumps(response_object), 404)

        # receive data
        try:
            content = request.get_json()
            file_list = content['file']
            filename_list = content['names']
            if len(file_list) == 0:
                response_object['message'] = 'Not file selected.'
                return make_response(json.dumps(response_object), 503)
        except Exception as error:
            response_object['message'] = 'Failed to read data because {}'.format(repr(error))
            return make_response(json.dumps(response_object), 404)

        # saving data
        os.makedirs(UPLOAD_DIR, exist_ok=True)
        os.makedirs('data/json', exist_ok=True)

        try:
            for i, file in enumerate(file_list):
                filename = os.path.join(UPLOAD_DIR, filename_list[i])
                encoded_text = file.split(',')[-1]
                text_buf = base64.b64decode(encoded_text)
                with open(filename, 'wb') as f:
                    f.write(text_buf)

                shutil.copy(filename, os.path.join('data', 'json', os.path.basename(filename)))
        except Exception as error:
            response_object['message'] = 'Failed to save data because {}'.format(repr(error))
            return make_response(json.dumps(response_object), 503)

        # call clustering module
        try:
            result = clustering(new_file_list=filename_list)
            response_object['status'] = 'success'
            response_object['message'] = '<br>'.join(["{}: {}".format(x[0], x[1]) for x in result])
        except Exception as error:
            response_object['message'] = repr(error)

        shutil.rmtree(UPLOAD_DIR, ignore_errors=True)
        
        return make_response(json.dumps(response_object, indent=2), 200)