"""Tiny byte-level CNN classifier for secret candidates.""" import torch import torch.nn as nn import torch.nn.functional as F from .encode import TYPES, VOCAB # noqa: F401 (TYPES re-exported for old imports) def pool_mask(m, L): """Exact adaptive_max_pool1d for a 0/1 mask (B, Lin) -> (B, L), written as a matmul so ONNX can export it.""" Lin, L = int(m.shape[-1]), int(L) # python ints: the matrix is a constant in the ONNX graph if Lin % L == 0: # equal windows: plain max pooling, no constant needed return F.max_pool1d(m.unsqueeze(1), Lin // L).squeeze(1) A = torch.zeros(Lin, L) for i in range(L): A[(i * Lin) // L:-(-((i + 1) * Lin) // L), i] = 1 return (m @ A).clamp(max=1) def _class_table(): """VOCAB x 8 one-hot byte classes: lower, upper, digit, hex letter, base64 symbol (+/=), other symbol, space, non-ASCII. Ids 0..2 (pad/separators) belong to no class.""" t = torch.zeros(VOCAB, 8) for v in range(3, VOCAB): c = chr(v - 3) if v - 3 >= 128: t[v, 7] = 1 elif c.islower(): t[v, 0] = 1 elif c.isupper(): t[v, 1] = 1 elif c.isdigit(): t[v, 2] = 1 elif c.isspace(): t[v, 6] = 1 else: t[v, 4 if c in "+/=" else 5] = 1 if c in "abcdefABCDEF": t[v, 3] = 1 return t N_FEATS = 13 CLASS_TABLE = _class_table() # built once at import, so ONNX export sees a constant def value_features(ids, seg): """Global statistics of the candidate bytes that a small CNN with max pooling sees poorly: length, Shannon entropy, character-class fractions and class switches (a memorability proxy, as in PassFinder / CredSweeper features). Computed inside the graph, so ONNX users need nothing new.""" m = ((seg == 2) & (ids > 2)).float() # candidate bytes, separators excluded n = m.sum(-1).clamp(min=1) cnt = torch.zeros(ids.shape[0], VOCAB).scatter_add(1, ids, m) # B, VOCAB byte counts p = cnt / n.unsqueeze(-1) ent = -(p * torch.log(p.clamp(min=1e-9))).sum(-1) cls = F.embedding(ids, CLASS_TABLE) * m.unsqueeze(-1) # B, L, 8 frac = cls.sum(1) / n.unsqueeze(-1) cid = cls.argmax(-1) both = m[:, 1:] * m[:, :-1] switches = ((cid[:, 1:] != cid[:, :-1]).float() * both).sum(-1) / n return torch.cat([(n / 96).unsqueeze(-1), (ent / 6.6).unsqueeze(-1), (ent / torch.log(n + 1)).unsqueeze(-1), frac, switches.unsqueeze(-1), (cnt.amax(-1) / n).unsqueeze(-1)], -1) class Block(nn.Module): def __init__(self, ch, k, dil): super().__init__() self.conv = nn.Conv1d(ch, ch, k, padding=dil * (k // 2), dilation=dil) self.norm = nn.BatchNorm1d(ch) def forward(self, x): return x + F.gelu(self.norm(self.conv(x))) class SecretCNN(nn.Module): def __init__(self, emb=32, ch=64, dilations=(1, 2, 4, 8), k=3, stride=1, n_types=len(TYPES), feats=False): super().__init__() self.config = dict(emb=emb, ch=ch, dilations=list(dilations), k=k, stride=stride, n_types=n_types) if feats: self.config["feats"] = True self.feats = feats self.stride = stride self.byte_emb = nn.Embedding(VOCAB, emb, padding_idx=0) self.seg_emb = nn.Embedding(4, emb) self.inp = nn.Conv1d(emb, ch, 3 if stride == 1 else 2 * stride, stride=stride, padding=1 if stride == 1 else stride // 2) self.blocks = nn.Sequential(*[Block(ch, k, d) for d in dilations]) self.head = nn.Sequential(nn.Linear(3 * ch + (N_FEATS if feats else 0), ch), nn.GELU(), nn.Dropout(0.1)) self.out = nn.Linear(ch, 1) self.type_out = nn.Linear(ch, n_types) def forward(self, ids, seg): x = self.byte_emb(ids) + self.seg_emb(seg) x = self.blocks(F.gelu(self.inp(x.transpose(1, 2)))) # B, C, L L = x.shape[-1] cand = pool_mask((seg == 2).float(), L).unsqueeze(1) anyt = pool_mask((seg > 0).float(), L).unsqueeze(1) neg = -1e4 cmax = (x + (1 - cand) * neg).amax(-1) cmean = (x * cand).sum(-1) / cand.sum(-1).clamp(min=1) amax = (x + (1 - anyt) * neg).amax(-1) parts = [cmax, cmean, amax] + ([value_features(ids, seg)] if self.feats else []) h = self.head(torch.cat(parts, -1)) return self.out(h).squeeze(-1), self.type_out(h) class Attn(nn.Module): """Plain multi-head self-attention (no fused kernels, so ONNX export stays simple).""" def __init__(self, d, heads): super().__init__() self.h, self.dh = heads, d // heads self.qkv = nn.Linear(d, 3 * d) self.proj = nn.Linear(d, d) def forward(self, x, mask): # x: B, L, D ; mask: B, L (1 = keep) # shape-free ops only, so the ONNX graph keeps a dynamic batch dimension B = x.shape[0] q, k, v = [t.reshape(B, -1, self.h, self.dh).transpose(1, 2) for t in self.qkv(x).chunk(3, -1)] att = (q @ k.transpose(-1, -2)) / (self.dh ** 0.5) att = att + (1 - mask.unsqueeze(1).unsqueeze(1)) * -1e4 y = (att.softmax(-1) @ v).transpose(1, 2).reshape(B, -1, self.h * self.dh) return self.proj(y) class TBlock(nn.Module): def __init__(self, d, heads, ff): super().__init__() self.n1, self.n2 = nn.LayerNorm(d), nn.LayerNorm(d) self.attn = Attn(d, heads) self.ff = nn.Sequential(nn.Linear(d, ff), nn.GELU(), nn.Linear(ff, d)) def forward(self, x, mask): x = x + self.attn(self.n1(x), mask) return x + self.ff(self.n2(x)) class SecretTransformer(nn.Module): """Byte transformer: conv stem (stride 2) -> N pre-norm transformer blocks -> same pooling head as the CNN.""" def __init__(self, emb=64, d=128, layers=2, heads=4, ff=256, stride=2, seq=None, n_types=len(TYPES), conv_dilations=()): super().__init__() from .encode import SEQ seq = seq or SEQ self.config = dict(arch="transformer", emb=emb, d=d, layers=layers, heads=heads, ff=ff, stride=stride, seq=seq, n_types=n_types, conv_dilations=list(conv_dilations)) # optional dilated conv blocks before attention ("hybrid"): local byte patterns first, then global context self.convs = nn.Sequential(*[Block(d, 3, dl) for dl in conv_dilations]) self.byte_emb = nn.Embedding(VOCAB, emb, padding_idx=0) self.seg_emb = nn.Embedding(4, emb) self.stem = nn.Conv1d(emb, d, 2 * stride, stride=stride, padding=stride // 2) L = (seq + 2 * (stride // 2) - 2 * stride) // stride + 1 self.pos = nn.Parameter(torch.zeros(1, L, d)) nn.init.normal_(self.pos, std=0.02) self.blocks = nn.ModuleList([TBlock(d, heads, ff) for _ in range(layers)]) self.norm = nn.LayerNorm(d) self.head = nn.Sequential(nn.Linear(3 * d, d), nn.GELU(), nn.Dropout(0.1)) self.out = nn.Linear(d, 1) self.type_out = nn.Linear(d, n_types) def forward(self, ids, seg): x = self.byte_emb(ids) + self.seg_emb(seg) x = self.convs(F.gelu(self.stem(x.transpose(1, 2)))).transpose(1, 2) # B, L, D L = x.shape[1] x = x + self.pos[:, :L] cand = pool_mask((seg == 2).float(), L) # B, L anyt = pool_mask((seg > 0).float(), L) for b in self.blocks: x = b(x, anyt) x = self.norm(x).transpose(1, 2) # B, D, L cand, anyt = cand.unsqueeze(1), anyt.unsqueeze(1) cmax = (x + (1 - cand) * -1e4).amax(-1) cmean = (x * cand).sum(-1) / cand.sum(-1).clamp(min=1) amax = (x + (1 - anyt) * -1e4).amax(-1) h = self.head(torch.cat([cmax, cmean, amax], -1)) return self.out(h).squeeze(-1), self.type_out(h) def build(config): config = dict(config) if config.pop("arch", "cnn") == "transformer": return SecretTransformer(**config) return SecretCNN(**config)