Spaces:
Sleeping
Sleeping
File size: 3,031 Bytes
e9d7935 6ffab9d e9d7935 d575357 e9d7935 d575357 e9d7935 d575357 e9d7935 d575357 e9d7935 d575357 e9d7935 6ffab9d d575357 e9d7935 d575357 e9d7935 d575357 e9d7935 d575357 e9d7935 d575357 e9d7935 d575357 e9d7935 6ffab9d d575357 e9d7935 d575357 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
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 pandas as pd
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_non_obfuscated.pkl", cpu=True) # change to "model_NonObf.pkl" if needed
model_Obf = load_learner("resnet18_obfuscated_samples.pkl", cpu=True) # change to "model_Shikata.pkl" if needed
def entropy_tester(bin_data):
entropy = stats.entropy(bin_data, base=2)
pred = entropy_classifier.predict(
pd.DataFrame(data=entropy.reshape(1, -1), columns=['Entropy'])
)
return pred[0]
def process_file(file):
# Gradio 4 with type="filepath" returns a string path.
# Older Gradio could return a file object with `.name`.
if file is None:
return "No file", "Please upload a file.", None, None
if isinstance(file, str):
file_path = file
else:
file_path = file.name
greyscale_data = get_binary_data(file_path)
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 == "Obfuscated":
prediction, _, probas = model_Obf.predict(converted_filename)
else:
# If some unexpected label appears, return a safe message
return (
"Unknown",
"Could not determine obfuscation type.",
None,
Image.open(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", type="filepath")
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, width=300)
button = gr.Button(value="Process File")
button.click(
process_file,
inputs=[file],
outputs=[text0, text1, prob_dist_img, converted_img],
)
demo.launch()
|