Spaces:
Sleeping
Sleeping
File size: 5,188 Bytes
88d577b 461e792 88d577b | 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | """
Inference script for Engine B (Malimg).
This module provides functionality to read raw binary files, convert them
to 2D grayscale image tensors, and classify their malware family using
a pre-trained Convolutional Neural Network (CNN).
"""
import os
import torch
from torchvision import transforms
from PIL import Image
import numpy as np
import json
from src.engine_b.model import MalwareCNN
import math
def bytes_to_image(file_path):
"""
Reads a raw binary file and converts its bytes into a square grayscale image.
If the file is a mock JSON profile, it reads the byte array from the JSON.
Otherwise, it reads the raw bytes of the executable and reshapes them into
a square 2D matrix, padding the end with zeros if necessary.
Args:
file_path (str): Path to the target binary file.
Returns:
PIL.Image: A grayscale (mode 'L') Image object representing the binary.
"""
# Check if this is our safe mock profile
if file_path.endswith(".json"):
try:
with open(file_path, "r") as f:
data = json.load(f)
if data.get("is_mock_profile"):
byte_array = np.array(data["malimg_bytes"], dtype=np.uint8)
else:
with open(file_path, "rb") as f:
binary_data = f.read()
byte_array = np.frombuffer(binary_data, dtype=np.uint8)
except:
with open(file_path, "rb") as f:
binary_data = f.read()
byte_array = np.frombuffer(binary_data, dtype=np.uint8)
else:
# Read raw binary
with open(file_path, "rb") as f:
binary_data = f.read()
byte_array = np.frombuffer(binary_data, dtype=np.uint8)
# Calculate image dimensions (square)
length = len(byte_array)
if length == 0:
return Image.new("L", (128, 128), color=0)
width = int(math.ceil(math.sqrt(length)))
height = width
# Pad array to form a perfect square
padded_length = width * height
padded_array = np.pad(byte_array, (0, padded_length - length), mode="constant")
# Reshape and create PIL Image
image_2d = padded_array.reshape((height, width))
img = Image.fromarray(image_2d, mode="L")
return img
class EngineBInfer:
"""
Inference Engine for Visual Malware Family Classification.
Attributes:
device (torch.device): CPU or CUDA device for inference.
classes (list): Ordered list of malware family class names.
model (MalwareCNN): The loaded PyTorch CNN.
transform (transforms.Compose): Image preprocessing pipeline.
"""
def __init__(
self,
model_path="models/engine_b_model.pth",
classes_path="models/engine_b_classes.json",
):
"""
Initializes the vision inference engine.
Args:
model_path (str): Path to the trained PyTorch state dictionary.
classes_path (str): Path to the JSON list mapping indices to class names.
"""
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.classes = []
if os.path.exists(classes_path):
with open(classes_path, "r") as f:
self.classes = json.load(f)
else:
# Fallback dummy classes if not trained yet
self.classes = [f"Class_{i}" for i in range(24)]
num_classes = len(self.classes)
self.model = MalwareCNN(num_classes=num_classes)
if os.path.exists(model_path):
self.model.load_state_dict(
torch.load(model_path, map_location=self.device, weights_only=True)
)
else:
print(f"Warning: {model_path} not found. Using untrained weights.")
self.model.to(self.device)
self.model.eval()
self.transform = transforms.Compose(
[
transforms.Resize((128, 128)),
transforms.ToTensor(),
]
)
def predict(self, file_path):
"""
Converts the target file to an image and runs CNN inference.
Args:
file_path (str): Path to the target file.
Returns:
dict: Contains 'family' (str), 'confidence' (float),
'all_probabilities' (dict), and 'image' (PIL.Image upscaled).
"""
img = bytes_to_image(file_path)
tensor = self.transform(img).unsqueeze(0).to(self.device)
with torch.no_grad():
outputs = self.model(tensor)
probabilities = torch.nn.functional.softmax(outputs, dim=1)[0]
# Create a dictionary of all class probabilities
all_probs = {
self.classes[i]: probabilities[i].item()
for i in range(len(self.classes))
}
top_prob, top_class = torch.max(probabilities, 0)
class_name = self.classes[top_class.item()]
return {
"family": class_name,
"confidence": top_prob.item(),
"all_probabilities": all_probs,
"image": img.resize((512, 512), resample=Image.NEAREST),
}
|