Shorter model card; rename package to password_detector

3e552ec verified 2 days ago

8 kB

"""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)

	"""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)