Spaces:
Sleeping
Sleeping
| 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() | |