web_deployment.py

"""Malware detection web application."""
import io
import base64
from os import getenv
import pathlib
import platform
import dash
from dash import dcc, html
from dash.dependencies import Input, Output, State
import plotly.express as px
from fastai.vision.all import load_learner
from binary2image import get_size, save_file

plt = platform.system()
if plt == "Linux":
    pathlib.WindowsPath = pathlib.PosixPath

learn = load_learner("model", cpu=True)
app = dash.Dash(__name__)
server = app.server

button_style = {
    "width": "20%",
    "height": "60px",
    "lineHeight": "60px",
    "borderWidth": "1px",
    "borderStyle": "dashed",
    "borderRadius": "5px",
    "textAlign": "center",
    "margin": "10px",
}
app.layout = html.Div(
    children=[
        html.H1(children="Malware detection"),
        html.Div(children="Select the executable file to classify:"),
        dcc.Upload(
            id="upload-data",
            children=html.Div(["Drag and Drop or ", html.A("Select a File")]),
            style=button_style,
        ),
        html.Br(),
        html.H2(id="description"),
        dcc.Graph(id="graph"),
    ]
)


@app.callback(
    Output(component_id="description", component_property="children"),
    Output(component_id="graph", component_property="figure"),
    [
        Input(component_id="upload-data", component_property="contents"),
        State(component_id="upload-data", component_property="filename"),
    ],
)
def update_message(contents, filename):
    """Show prediction result to user."""
    if contents is not None:
        _, content_string = contents.split(",")
        decoded = base64.b64decode(content_string)
        file = io.BytesIO(decoded)
        greyscale_data = []
        while byte := file.read(1):
            greyscale_data.append(ord(byte))
        bin_size = get_size(len(greyscale_data))
        save_file(".", filename, greyscale_data, bin_size)
        prediction, decoded, probas = learn.predict(filename)
        message = f"Your file is {prediction}!!!"
        fig = px.pie(probas, values=probas, names=["Goodware", "Malware"])
        return message, fig


if __name__ == "__main__":
    app.run_server(host="0.0.0.0", port=getenv("PORT", "80"))