final

src/config/config.py

@@ -0,0 +1,18 @@
+from transformers import BertTokenizer
+import torch
+
+
+configuration = {
+    "sequence_length": 100,
+    "tokenizer_name": "bert-base-uncased",
+    "vocab_size": len(BertTokenizer.from_pretrained("bert-base-uncased")),
+    "embedding_dim": 10,
+    "input_length": 100,
+    "num_filters": 128,
+    "kernel_size": 4,
+    "num_gated_units": 64,
+    "hidden_neurons": 128,
+    "dropout_cnn": 0.2,
+    "dropout_fc": 0.5,
+    "device": "cuda"
+}

src/datasets/datasets.py

@@ -0,0 +1,41 @@
+import torch
+from torch.utils.data import Dataset
+
+from transformers import BertTokenizer
+from PIL import Image
+import numpy as np
+
+from typing import List
+
+
+class CombinedDataset(Dataset):
+    
+    def __init__(self, api_call_list, img_path, sequence_length, max_len=128, transforms=None, tokenizer_name='bert-base-uncased'):
+    
+        self.image_path = img_path
+
+        self.transforms = transforms
+        self.max_len = max_len
+        self.sequence_length = sequence_length
+        
+
+        self.tokenizer = BertTokenizer.from_pretrained(tokenizer_name)
+        self.api_calls = api_call_list
+        self.encoded_calls = [self.tokenizer.encode(" ".join(call), add_special_tokens=True, max_length=self.max_len, padding='max_length', truncation=True) for call in self.api_calls]
+        self.padded_calls = np.array([x + [0] * (self.max_len - len(x)) if len(x) < self.max_len else x[:self.max_len] for x in self.encoded_calls])
+        print("Dataset initialized")
+
+    def __len__(self):
+        return len(self.padded_calls)
+
+    def __getitem__(self,idx):
+        img_path = self.image_path
+        image = Image.open(img_path)
+        
+        if self.transforms:
+            image = self.transforms(image)
+
+        tokenized_seq = self.padded_calls
+        
+
+        return torch.tensor(tokenized_seq, dtype=torch.long), image

src/main/make_predictions.py

@@ -0,0 +1,55 @@
+from src.config import config
+from src.datasets.datasets import CombinedDataset
+from src.models.multimodal import CombinedMalwareDetectionModel
+from src.models.bigru import CNNBiGRU
+from src.models.cnn import ImprovedCNN
+from src.utils.get_features import get_img_api
+from torchvision import transforms
+
+import pickle
+import torch
+
+
+data_path = 'src/data/subset_dataset.csv'
+
+simple_transform = transforms.Compose([
+    transforms.Resize((128, 128)),
+    transforms.RandomHorizontalFlip(p=0.5),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.5], std=[0.5])  
+])
+
+def load_model(model_path, device='cpu'):
+    """Loads the model from a pickle file and moves it to the specified device."""
+    with open(model_path, 'rb') as f:
+        model = pickle.load(f)
+    return model.to(device)
+
+
+def get_prediction(model, padded_sequences, img_x, device='cuda'):
+    malware_classes = ["Benign", "RedLine Stealer", "Downloader", "RAT", 
+                       "Banking Trojan", "Snake Keylogger", "Spyware"]
+
+    # Move inputs to the device
+    padded_sequences, img_x = padded_sequences.to(device), img_x.to(device)
+    
+    # Perform inference
+    outputs = model(padded_sequences, img_x)
+    _, predicted = torch.max(outputs, 1)  
+    
+    return malware_classes[predicted]
+
+
+model_path = "model_dump/model_malware_lstm (1).pkl"
+
+image_path, api_call_list = get_img_api('002ce0d28ec990aadbbc89df457189de37d8adaadc9c084b78eb7be9a9820c81', data_path)
+
+
+dataset = CombinedDataset(api_call_list, image_path, transforms=simple_transform ,sequence_length=config.configuration["sequence_length"])
+
+padded_sequences,img_x=next(iter(dataset))
+img_x=img_x.unsqueeze(0) #type: ignore
+
+model = load_model(model_path, device=config.configuration["device"])
+predicted_class = get_prediction(model, padded_sequences, img_x, config.configuration["device"])
+print(f"Predicted class: {predicted_class}")

src/models/bigru.py

@@ -0,0 +1,32 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class CNNBiGRU(nn.Module):
+    def __init__(self, vocab_size, embedding_dim, input_length, num_filters, kernel_size, 
+                 num_gated_units, hidden_neurons, dropout_cnn, dropout_fc):
+        super(CNNBiGRU, self).__init__()
+        
+        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
+        self.conv1d = nn.Conv1d(in_channels=embedding_dim, out_channels=num_filters, kernel_size=kernel_size)
+        self.dropout_cnn = nn.Dropout(dropout_cnn)
+        self.maxpool = nn.MaxPool1d(kernel_size=kernel_size, stride=1)
+        self.bigru = nn.LSTM(input_size=num_filters, hidden_size=num_gated_units, num_layers=1, 
+                             batch_first=True, bidirectional=True)
+        
+        self.fc1 = nn.Linear(num_gated_units * 2, hidden_neurons) 
+        self.fc2 = nn.Linear(hidden_neurons, hidden_neurons)
+        self.dropout_fc = nn.Dropout(dropout_fc)
+        self.output = nn.Linear(hidden_neurons, 128)  
+        
+    def forward(self, x):
+        x = self.embedding(x)  
+        x = x.permute(0, 2, 1)  
+        x = F.relu(self.conv1d(x))
+        x = self.dropout_cnn(x)
+        x = self.maxpool(x)  
+        x = x.permute(0, 2, 1)  
+        x, _ = self.bigru(x)  
+        x = x[:, -1, :]  
+        x = self.output(x)
+        return x

src/models/cnn.py

@@ -0,0 +1,41 @@
+import torch
+import torch.nn as nn
+
+class ImprovedCNN(nn.Module):
+    def __init__(self, input_channels, hidden_units, num_classes=4):
+        super().__init__()
+        self.block1 = nn.Sequential(
+            nn.Conv2d(in_channels=input_channels, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(hidden_units),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2)
+        )
+        
+        self.block2 = nn.Sequential(
+            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units*2, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(hidden_units*2),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2)
+        )
+        
+        self.block3 = nn.Sequential(
+            nn.Conv2d(in_channels=hidden_units*2, out_channels=hidden_units*4, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(hidden_units*4),
+            nn.ReLU(),
+            nn.AdaptiveAvgPool2d(output_size=(4, 4))
+        )
+        
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(hidden_units*4*4*4, 256),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(256, 128)
+        )
+    
+    def forward(self, x):
+        x = self.block1(x)
+        x = self.block2(x)
+        x = self.block3(x)
+        x = self.classifier(x)
+        return x

src/models/multimodal.py

@@ -0,0 +1,36 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from src.models.bigru import CNNBiGRU
+from src.models.cnn import ImprovedCNN
+from src.config import config
+
+
+class CombinedMalwareDetectionModel(nn.Module):
+    def __init__(self, vocab_size, embedding_dim, num_filters, kernel_size):
+        super(CombinedMalwareDetectionModel, self).__init__()
+
+        self.malware_detection_model = CNNBiGRU(vocab_size, embedding_dim,config.configuration["input_length"], num_filters, kernel_size, 
+                                                config.configuration["num_gated_units"], config.configuration["hidden_neurons"],
+                                                config.configuration["dropout_cnn"], config.configuration["dropout_fc"]) 
+        
+        self.improved_cnn = ImprovedCNN(input_channels=1, hidden_units=32)
+
+        self.fc1 = nn.Linear(256, 64)  
+        self.fc2 = nn.Linear(64, 32)
+        self.fc3 = nn.Linear(32, 7)  
+
+        self.dropout = nn.Dropout(0.2)
+
+    def forward(self, padded_sequences, img_x):
+        output_api = self.malware_detection_model(padded_sequences)
+        output_img = self.improved_cnn(img_x)
+
+        input_multi = torch.cat([output_img, output_api], dim=-1).to(torch.float32) 
+        
+        x = F.relu(self.fc1(input_multi)) 
+        x = self.dropout(x)
+        x = F.relu(self.fc2(x))
+        x = self.dropout(x)
+        x = self.fc3(x) 
+        return x

src/utils/get_features.py

@@ -0,0 +1,31 @@
+import pandas as pd
+
+def get_img_api(hash, data_path):
+ 
+    df = pd.read_csv(data_path)
+    
+    
+    row = df[df['SHA256'] == hash]
+    
+
+    if row.empty:
+        return None, None
+    
+    # Extract the image path
+    img_path = 'src/data/images/' + row['SHA256'].values[0] + '.jpg'
+    
+    # Extract the API calls
+    api_columns = df.columns[2:]  # Skip the first two columns (SHA256 and Type)
+    api_calls = row[api_columns].values.flatten().tolist()
+    
+    # Filter out only the API calls that are present (value == 1)
+    api_call_list = [[api for api, value in zip(api_columns, api_calls) if value == 1]]
+    
+    return img_path, api_call_list
+
+
+# hash_value = '002ce0d28ec990aadbbc89df457189de37d8adaadc9c084b78eb7be9a9820c81'
+# img_path, api_call_list = get_img_api(hash_value)
+
+# print("Image Path:", img_path)
+# print("API Call List:", api_call_list)