| """SecretDetector: find secrets and credentials in code, configs, logs, chats and documents with a compact Ettin |
| token tagger. |
| |
| from cbx_secrets import SecretDetector |
| det = SecretDetector() # this folder: model.safetensors, config.json, tokenizer, secrets_config.json |
| det.extract('export STRIPE_KEY="sk_live_..."') |
| det.redact('export STRIPE_KEY="sk_live_..."') |
| |
| (Not named secrets.py: that would hide Python's standard `secrets` module.) |
| |
| The model tags every token with a BIO tag (B-/I- per label, O = not a secret); tags are joined into character spans. |
| A text is read in one pass, up to the model's full context length (7,999 tokens, from config.json). Only a text |
| longer than that is read in overlapping pieces; a token in an overlap takes its tag from the piece where it sits |
| further from the edge, so a value split by a piece boundary is still read in one piece. |
| """ |
| import json, os, time |
|
|
| import torch |
| import torch.nn as nn |
| from safetensors.torch import load_file |
| from transformers import AutoConfig, AutoModel, AutoTokenizer |
|
|
| _EDGE_LEAD, _EDGE_TRAIL = "=.\"'*<:;,`", ".,;:\"'*>`" |
|
|
|
|
| class _Tagger(nn.Module): |
| def __init__(self, config, n_tags): |
| super().__init__() |
| self.backbone = AutoModel.from_config(config, attn_implementation="sdpa") |
| self.head = nn.Linear(config.hidden_size, n_tags) |
|
|
| def forward(self, input_ids): |
| h = self.backbone(input_ids=input_ids, attention_mask=torch.ones_like(input_ids)).last_hidden_state |
| return self.head(h.float()) |
|
|
|
|
| class SecretDetector: |
| def __init__(self, path=None, device=None, overlap=64, join="none"): |
| path = path or os.path.dirname(os.path.abspath(__file__)) |
| self.cfg = json.load(open(os.path.join(path, "secrets_config.json"))) |
| self.tags, self.labels = self.cfg["tags"], self.cfg["labels"] |
| self.overlap, self.join = overlap, join |
| self.min_len = self.cfg.get("min_len") or {} |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| self.tok = AutoTokenizer.from_pretrained(path) |
| config = AutoConfig.from_pretrained(path) |
| self.max_length = config.max_position_embeddings |
| self.model = _Tagger(config, len(self.tags)) |
| self.model.load_state_dict(load_file(os.path.join(path, "model.safetensors"))) |
| self.model.to(self.device).eval() |
|
|
| |
| @torch.inference_mode() |
| def _tag(self, ids): |
| """One tag id per token: the whole text in one pass, or overlapping pieces if it exceeds the context.""" |
| step = self.max_length - 2 |
| if len(ids) <= step: starts = [0] |
| else: |
| stride = step - self.overlap |
| starts = list(range(0, len(ids) - step, stride)) + [len(ids) - step] |
| out = [0] * len(ids) |
| for k, s in enumerate(starts): |
| e = min(s + step, len(ids)) |
| x = torch.tensor([[self.tok.cls_token_id] + ids[s:e] + [self.tok.sep_token_id]], device=self.device) |
| with torch.autocast(device_type="cuda", dtype=torch.float16, enabled=self.device == "cuda"): |
| pred = self.model(x).argmax(-1)[0, 1:-1].tolist() |
| lo = s if k == 0 else s + self.overlap // 2 |
| hi = e if k == len(starts) - 1 else e - self.overlap // 2 |
| for i in range(lo, hi): out[i] = pred[i - s] |
| return out |
|
|
| def _spans(self, text, offsets, tag_ids): |
| spans, cur = [], None |
| for (a, b), t in zip(offsets, tag_ids): |
| while a < b and text[a].isspace(): a += 1 |
| while b > a and text[b - 1].isspace(): b -= 1 |
| if a >= b: continue |
| tag = self.tags[t] |
| if tag == "O": |
| if cur: spans.append(cur); cur = None |
| continue |
| prefix, label = tag.split("-", 1) |
| if prefix == "B" or cur is None or cur[2] != label: |
| if cur: spans.append(cur) |
| cur = [a, b, label] |
| else: |
| cur[1] = b |
| if cur: spans.append(cur) |
| out = [] |
| for a, b, label in spans: |
| while a < b and text[a] in _EDGE_LEAD: a += 1 |
| while b > a and text[b - 1] in _EDGE_TRAIL: b -= 1 |
| if b > a: out.append(dict(label=label, text=text[a:b], start=a, end=b)) |
| out = self._join(text, out) if self.join == "word" else out |
| return [s for s in out if s["end"] - s["start"] >= self.min_len.get(s["label"], self.min_len.get("*", 1))] |
|
|
| @staticmethod |
| def _join(text, spans): |
| """Optional: spans separated only by non-whitespace become one span, labelled by the label covering most of its |
| characters. Off by default: it merges adjacent secrets such as `user:password`.""" |
| out = [] |
| for s in spans: |
| gap = text[out[-1]["end"]:s["start"]] if out else " " |
| if out and len(gap) <= 20 and not any(c.isspace() for c in gap): |
| p = out[-1] |
| p["chars"][s["label"]] = p["chars"].get(s["label"], 0) + s["end"] - s["start"] |
| p.update(end=s["end"], text=text[p["start"]:s["end"]], label=max(p["chars"], key=p["chars"].get)) |
| else: |
| out.append(dict(s, chars={s["label"]: s["end"] - s["start"]})) |
| for s in out: del s["chars"] |
| return out |
|
|
| |
| def extract(self, text, labels=None): |
| """List of {label, text, start, end}; `labels` keeps only those types.""" |
| enc = self.tok(text, add_special_tokens=False, return_offsets_mapping=True, verbose=False) |
| spans = self._spans(text, enc["offset_mapping"], self._tag(enc["input_ids"])) |
| return [s for s in spans if labels is None or s["label"] in labels] |
|
|
| def extract_batch(self, texts, labels=None): |
| return [self.extract(t, labels) for t in texts] |
|
|
| def redact(self, text, labels=None, template="[{label}]"): |
| """The text with every found secret replaced by `template` (default "[API_KEY]", "[PASSWORD]", ...).""" |
| out, last = [], 0 |
| for s in self.extract(text, labels): |
| out += [text[last:s["start"]], template.format(label=s["label"])]; last = s["end"] |
| return "".join(out + [text[last:]]) |
|
|
| def timed_extract(self, text): |
| t = time.perf_counter(); spans = self.extract(text) |
| return spans, (time.perf_counter() - t) * 1000 |
|
|
