[原创]KCTF 2026 第九题 "星海抉择" WriteUp-CTF对抗-看雪安全社区

## 1. 题目信息

题目给了三个关键文件：

- `model_def.py`：模型和 tokenizer 的定义。
- `inference.py`：加载模型并对输入预测一个字符。
- `ictf_model/model.safetensors`：模型权重。

题目目标是分析模型代码和权重，找出隐藏在模型里的秘密，也就是能让模型输出成功标记的输入串。

## 2. 先看推理流程

`inference.py` 的核心逻辑很短：

```python
user_input = input("Enter your prompt: ")
inputs = tokenizer(user_input, return_tensors="pt")
input_ids = inputs["input_ids"]

with torch.no_grad():
    outputs = model(input_ids)
    logits = outputs["logits"]
    next_token_id = torch.argmax(logits[:, -1, :], dim=-1).item()
    predicted_char = tokenizer.id2char.get(next_token_id, "<unk>")
    print(predicted_char)
```

它不会生成一整段文本，只会对输入字符串预测一个输出字符。

所以问题可以简化为：

> 找一个输入字符串，让模型最后 `argmax` 出来的类别是某个特殊字符。

再看 `ICTFTokenizer`：

```python
self.charset = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
self.char2id = {c: i for i, c in enumerate(self.charset)}
self.id2char = {i: c for i, c in enumerate(self.charset)}
self.id2char[62] = "<success>"
self.id2char[63] = "<fail>"
self.vocab_size = 64
self.pad_token_id = 0
```

输入字符只能来自：

```text
0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ
```

也就是 62 个普通字符。

模型额外定义了两个输出类别：

- `62`：`<success>`
- `63`：`<fail>`

因此正确方向很明确：找一个长度不超过 16 的字符串，使模型输出 `<success>`。

## 3. 模型结构非常简单

`model_def.py` 里的模型结构是：

```python
class ICTFForCausalLM(PreTrainedModel):
    config_class = ICTFConfig

def __init__(self, config):
        super().__init__(config)
        self.dense = nn.Linear(config.max_position_embeddings, config.hidden_size, bias=True)
        self.act = nn.ReLU()
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=True)

def forward(self, input_ids, **kwargs):
        x = input_ids.float()
        hidden_states = self.dense(x)
        hidden_states = self.act(hidden_states)
        logits = self.lm_head(hidden_states)
        logits = logits.unsqueeze(1)
        return {"logits": logits}
```

配置中：

```text
max_position_embeddings = 16
hidden_size = 21
vocab_size = 64
```

所以模型其实就是：

```text
16 维输入
  -> Linear(16, 21)
  -> ReLU
  -> Linear(21, 64)
  -> 取最大 logit
```

这里没有真正的 Transformer，也没有上下文生成逻辑。输入字符串会被 tokenizer 转成 16 个字符 ID，不足 16 位的地方补 `0`。

## 4. 读取权重，定位异常输出类

由于 `model.safetensors` 只是保存了几个张量，直接读取后可以看到：

```text
dense.bias      shape = [21]
dense.weight    shape = [21, 16]
lm_head.bias    shape = [64]
lm_head.weight  shape = [64, 21]
```

关键点在 `lm_head`。

检查 `lm_head.weight` 和 `lm_head.bias` 后发现，普通字符类别几乎全部被压低：

```text
class 0..61 bias = -10000
```

真正参与竞争的只有：

```text
class 62 <success>
class 63 <fail>
```

其中：

```text
<fail>:
  logit = 0.4

<success>:
  bias = -376131.21875
  weight[0] = 1
  weight[1..20] = -10000000000
```

也就是说，设 ReLU 后的隐藏层为：

```text
h0, h1, h2, ..., h20
```

则：

```text
fail_logit = 0.4

success_logit =
    h0
    - 10000000000 * (h1 + h2 + ... + h20)
    - 376131.21875
```

这个设计很明显：

1. 如果 `h1` 到 `h20` 里面任意一个大于 0，`success_logit` 会被 `-10000000000` 直接打爆，必然失败。
2. 如果 `h1` 到 `h20` 全部等于 0，那么 `success_logit = h0 - 376131.21875`。
3. 想让 `<success>` 赢过 `<fail>`，需要：

```text
success_logit > fail_logit
h0 - 376131.21875 > 0.4
h0 > 376131.61875
```

所以目标变成了一个非常清楚的约束问题：

```text
h1..h20 全部 <= 0，经 ReLU 后变成 0
h0 > 376131.61875
```

## 5. 把模型转成整数线性约束

输入是 16 个字符 ID：

```text
x0, x1, ..., x15
```

每个字符 ID 范围是：

```text
0 <= xi <= 61
```

`dense` 层是线性变换：

```text
pre_j = dense.weight[j] · x + dense.bias[j]
h_j = ReLU(pre_j)
```

因此成功条件可以写成：

```text
pre_1 <= 0
pre_2 <= 0
...
pre_20 <= 0

pre_0 > 376131.61875
```

因为权重大部分是整数，`pre_0` 的 bias 是 `311558.71875`，所以：

```text
311558.71875 + dense.weight[0] · x > 376131.61875
```

等价于：

```text
dense.weight[0] · x >= 64573
```

这样就不需要暴力跑 `62^16` 的输入空间了，直接用约束求解器求整数解即可。

## 6. 求解脚本

下面是完整的求解思路，直接解析 `safetensors`，不依赖 PyTorch：

```python
import struct
import json
import numpy as np
from z3 import *

charset = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

data = open(r"ictf_model/model.safetensors", "rb").read()
header_len = struct.unpack("<Q", data[:8])[0]
header = json.loads(data[8:8 + header_len])
offset = 8 + header_len

tensors = {}
for name, meta in header.items():
    if name == "__metadata__":
        continue

start, end = meta["data_offsets"]
    arr = np.frombuffer(data[offset + start:offset + end], dtype="<f4")
    tensors[name] = arr.reshape(meta["shape"]).copy()

W1 = tensors["dense.weight"].astype(int)
b1 = tensors["dense.bias"]

xs = [Int(f"x{i}") for i in range(16)]
s = Solver()

for x in xs:
    s.add(x >= 0, x <= 61)

# h1..h20 必须全部不激活
for j in range(1, 21):
    s.add(
        Sum([int(W1[j, i]) * xs[i] for i in range(16)])
        + int(round(float(b1[j])))
        <= 0
    )

# h0 需要超过 success 阈值
s.add(Sum([int(W1[0, i]) * xs[i] for i in range(16)]) >= 64573)

print(s.check())

m = s.model()
vals = [m[x].as_long() for x in xs]
print(vals)
print("".join(charset[v] for v in vals))
```

输出：

```text
sat
[15, 1, 10, 16, 2, 0, 2, 6, 12, 7, 15, 10, 1, 6, 6, 6]
f1ag2026c7fa1666
```

根据 charset 映射：

```text
15 -> f
1  -> 1
10 -> a
16 -> g
2  -> 2
0  -> 0
2  -> 2
6  -> 6
12 -> c
7  -> 7
15 -> f
10 -> a
1  -> 1
6  -> 6
6  -> 6
6  -> 6
```

得到：

```text
f1ag2026c7fa1666
```

## 7. 验证

把这组输入代回第一层，得到：

```text
pre_0 = 376131.71875
pre_1..pre_20 = 0
```

经过 ReLU 后：

```text
h0 = 376131.71875
h1..h20 = 0
```

所以：

```text
success_logit = 376131.71875 - 376131.21875 = 0.5
fail_logit = 0.4
```

`<success>` 的 logit 更大，因此模型输出：

```text
<success>
```

同时再额外加一条约束，要求解不能等于已找到的这组输入：

```python
s.add(Or([xs[i] != known[i] for i in range(16)]))
```

求解结果为：

```text
unsat
```

说明在当前约束下，这个成功输入是唯一的。

## 8. 总结

这题的关键不是把模型当成大语言模型去猜 prompt，而是先看清楚模型结构：

- tokenizer 把输入限制成 16 个字符 ID。
- 模型只是两层全连接加 ReLU。
- 输出层只保留 `<success>` 和 `<fail>` 两个有效类别。
- `<success>` 被设计成一个线性约束检查器：
  - `h1..h20` 必须全部不激活。
  - `h0` 必须超过指定阈值。

所以最终可以把 AI 模型逆向成整数线性约束问题，用 Z3 直接求解。

最终 flag / 注册码：

```text
f1ag2026c7fa1666
```

冰与火的战歌：Windows内核攻防实战高级班！从零到实战，融合AI与Windows内核攻防全技术栈，打造具备自动化能力的内核开发高手。