Spaces:
Sleeping
Sleeping
File size: 2,544 Bytes
e9d7935 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
# import io
# import base64
from PIL import Image
from fastai.vision.all import load_learner
from binary2image import get_size, save_file, get_binary_data
import gradio as gr
import numpy as np
import plotly.express as px
from scipy import stats
import pickle
## Loading the models
entropy_classifier = pickle.load(open('entropy_tester_classifier.pkl', 'rb'))
model_NonObf = load_learner("model.pkl", cpu=True) # change to "model_NonObf.pkl"
model_Shikata = load_learner("model.pkl", cpu=True) # change to "model_Shikata.pkl"
model_XOR = load_learner("model.pkl", cpu=True) # change to "model_XOR.pkl"
def entropy_tester(bin_data):
entropy = stats.entropy(bin_data,base=2)
pred = entropy_classifier.predict(
np.array(entropy).reshape(1, -1))
return pred[0]
def process_file(file):
greyscale_data = get_binary_data(file.name)
pred_entropy = entropy_tester(greyscale_data)
bin_size = get_size(len(greyscale_data))
save_file(".", "tempfile", greyscale_data, bin_size)
converted_filename = "tempfile.png"
# Make prediction depending on the type of obfuscation
if pred_entropy == "NonObf":
prediction, _, probas = model_NonObf.predict(converted_filename)
elif pred_entropy == "Shikata":
prediction, _, probas = model_Shikata.predict(converted_filename)
elif pred_entropy == "XOR":
prediction, _, probas = model_XOR.predict(converted_filename)
message = f"Your file is {prediction}!!!"
if pred_entropy == "NonObf": pred_entropy = "Non-Obfuscated"
# Convert probas to percentages
probas_percentage = [float(prob) * 100 for prob in probas]
fig = px.bar(x=["Goodware", "Malware"],
y=probas_percentage, labels={'x':'Type', 'y':'Probability (%)'},
height=300)
return pred_entropy, message, fig, Image.open("tempfile.png")
# Define the layout using Blocks, Row, and Column
with gr.Blocks() as demo:
with gr.Column(scale=1):
file = gr.File(label="Upload Executable File")
with gr.Row():
with gr.Column():
text1 = gr.Textbox(label="Prediction")
prob_dist_img = gr.Plot(label="Probability Distribution")
with gr.Column():
text0 = gr.Textbox(label="Type of Obfuscation")
converted_img = gr.Image(label="Converted Image", height=300)
button = gr.Button(value="Process File")
button.click(process_file, inputs=[file], outputs=[text0, text1, prob_dist_img, converted_img])
demo.launch()
|