Static-Malware-Classification/malware_web_app.py at main · tharunsridhar/Static-Malware-Classification · GitHub

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
# -*- coding: utf-8 -*-
"""app.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1vhvbHlK-xKlTTbGP9AqNOD-wPhJQR8E_
"""

import streamlit as st
import tensorflow as tf
import numpy as np
from PIL import Image
import tempfile
from tensorflow.keras.applications.efficientnet_v2 import preprocess_input

# ---------------- CONFIG ----------------
IMG_SIZE = 384
THRESHOLD = 0.85

DETECTION_MODEL = "/content/drive/MyDrive/Malware classifications/(detection)efficientnet v2s.h5"
CLASSIFICATION_MODEL = "/content/drive/MyDrive/Malware classifications/(classification)efficientnet v2s.h5"

CLASS_NAMES = [
    "Adposhel","Agent","Allaple","Alueron.gen!J","Amonetize","Androm",
    "Autorun","BrowseFox","C2LOP.gen!g","Dialplatform.B","Dinwod","Elex",
    "Expiro","Fakerean","Fasong","HackKMS","Hlux","Injector","InstallCore",
    "Lolyda.AA1","Lolyda.AA2","MultiPlug","Neoreklami","Neshta","Regrun",
    "Sality","Snarasite","Stantinko","VBA","VBKrypt","Vilsel"
]

st.set_page_config(page_title="Malware Detection", layout="centered")
st.title("🛡️ Malware Detection System")
st.caption("Static analysis only • Using .h5 models")

# ---------------- LOAD MODELS ----------------
@st.cache_resource
def load_models():
    det = tf.keras.models.load_model(DETECTION_MODEL, compile=False)
    cls = tf.keras.models.load_model(CLASSIFICATION_MODEL, compile=False)
    return det, cls

det_model, cls_model = load_models()

# ---------------- EXE → IMAGE ----------------
def exe_to_image(path, width=256):
    with open(path, "rb") as f:
        data = np.frombuffer(f.read(), dtype=np.uint8)
    h = int(np.ceil(len(data) / width))
    data = np.pad(data, (0, width*h - len(data)))
    img = data.reshape(h, width)
    img = Image.fromarray(img).convert("L")
    img = img.resize((IMG_SIZE, IMG_SIZE))
    return img.convert("RGB")

# ---------------- ANALYZE ----------------
def analyze(uploaded):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".exe") as tmp:
        tmp.write(uploaded.read())
        path = tmp.name

    img = exe_to_image(path)
    x = preprocess_input(np.array(img))
    x = np.expand_dims(x, axis=0)

    # Model 1: Detection
    det_pred = det_model.predict(x, verbose=0)[0]
    malware_prob = float(det_pred[1])

    if malware_prob < THRESHOLD:
        return "BENIGN", (1 - malware_prob) * 100, None

    # Model 2: Classification
    cls_pred = cls_model.predict(x, verbose=0)[0]
    idx = int(np.argmax(cls_pred))
    return "MALWARE", malware_prob * 100, CLASS_NAMES[idx]

# ---------------- UI ----------------
uploaded = st.file_uploader("Upload EXE file", type=["exe"])

if uploaded and st.button("Analyze"):
    with st.spinner("Analyzing..."):
        status, conf, family = analyze(uploaded)

    st.subheader(f"Result: {status}")
    st.write(f"Confidence: {conf:.2f}%")
    if family:
        st.write(f"Malware Family: {family}")
else:
    st.info("Upload an EXE file to get results.")