first commit

app_gradio.py

@@ -0,0 +1,72 @@
+# 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()

binary2image.py

@@ -0,0 +1,91 @@
+"""Binary to Image Converter."""
+import os
+from argparse import ArgumentParser
+from PIL import Image
+from tqdm import tqdm
+
+
+def get_binary_data(filename):
+    """Extract byte values from binary executable file and store them in list.
+
+    :param file: executable file
+    :return: byte value list
+    """
+    binary_values = []
+    with open(filename, "rb") as file:
+        while byte := file.read(1):
+            binary_values.append(ord(byte))
+
+    return binary_values
+
+
+def get_size(data_length):
+    """Obtain image size.
+
+    Source: Malware images: visualization and automatic classification.
+    :param data_length: Number of bytes in file
+    :return: size as integer tuple
+    """
+    size = data_length
+    kib = 2**10
+
+    if size < 10 * kib:
+        width = 32
+    elif size < 30 * kib:
+        width = 64
+    elif size < 60 * kib:
+        width = 128
+    elif size < 100 * kib:
+        width = 256
+    elif size < 200 * kib:
+        width = 384
+    elif size < 500 * kib:
+        width = 512
+    elif size < 1000 * kib:
+        width = 768
+    else:
+        width = 1024
+
+    height = size // width + 1
+
+    return (width, height)
+
+
+def save_file(folder, filename, data, size):
+    """Save PIL image to disk.
+
+    :param folder: folder where images will be saved
+    :param filename: binary filename
+    :param data: grayscale image
+    :param size: image size
+    """
+    image = Image.new("L", size)
+    image.putdata(data)
+
+    name, _ = os.path.splitext(filename)
+    name = os.path.basename(name)
+    imagename = os.path.join(folder, name + ".png")
+
+    image.save(imagename)
+
+
+if __name__ == "__main__":
+
+    parser = ArgumentParser(
+        description="Transform all files in a folder into PNG images"
+    )
+    parser.add_argument("input_folder", help="Folder with the original files")
+    parser.add_argument(
+        "output_folder", help="Folder where the images will be saved"
+    )
+    args = parser.parse_args()
+
+    files = [
+        os.path.join(args.input_folder, bin_file)
+        for bin_file in os.listdir(args.input_folder)
+    ]
+    os.makedirs(args.output_folder, exist_ok=True)
+    for bin_file in tqdm(files):
+        greyscale_data = get_binary_data(bin_file)
+        bin_size = get_size(len(greyscale_data))
+        save_file(args.output_folder, bin_file, greyscale_data, bin_size)

entropy_tester_classifier.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dbb214565a0e5969901fb4707f41369040ab882d8b72531060c404e2d1cc4c80
+size 1895142

model.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e3165c874eb66655b5e859c8733feae933fa0617eb988ec7be91790bf4e2e2f3
+size 47462689

requirements.txt

@@ -0,0 +1,8 @@
+# CPU-only version of PyTorch
+-f https://download.pytorch.org/whl/cpu/torch_stable.html
+torch==2.0.1
+torchvision==0.15.2
+
+# fastai library
+fastai==2.7.12
+timm