code-airlock 1.0.0__tar.gz → 1.3.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {code_airlock-1.0.0/src/code_airlock.egg-info → code_airlock-1.3.0}/PKG-INFO +97 -10
- {code_airlock-1.0.0 → code_airlock-1.3.0}/README.md +96 -9
- {code_airlock-1.0.0 → code_airlock-1.3.0}/pyproject.toml +1 -1
- {code_airlock-1.0.0 → code_airlock-1.3.0}/src/code_airlock/__init__.py +1 -1
- code_airlock-1.3.0/src/code_airlock/__main__.py +127 -0
- code_airlock-1.3.0/src/code_airlock/config.py +39 -0
- {code_airlock-1.0.0 → code_airlock-1.3.0}/src/code_airlock/engine.py +36 -16
- code_airlock-1.3.0/src/code_airlock/executor.py +87 -0
- code_airlock-1.3.0/src/code_airlock/mcp_server.py +124 -0
- {code_airlock-1.0.0 → code_airlock-1.3.0}/src/code_airlock/proxy.py +1 -1
- {code_airlock-1.0.0 → code_airlock-1.3.0}/src/code_airlock/sanitizer.py +27 -7
- code_airlock-1.3.0/src/code_airlock/scanner.py +76 -0
- {code_airlock-1.0.0 → code_airlock-1.3.0/src/code_airlock.egg-info}/PKG-INFO +97 -10
- {code_airlock-1.0.0 → code_airlock-1.3.0}/src/code_airlock.egg-info/SOURCES.txt +3 -0
- code_airlock-1.3.0/tests/test_mcp.py +181 -0
- code_airlock-1.0.0/src/code_airlock/__main__.py +0 -6
- code_airlock-1.0.0/src/code_airlock/mcp_server.py +0 -71
- code_airlock-1.0.0/tests/test_mcp.py +0 -66
- {code_airlock-1.0.0 → code_airlock-1.3.0}/LICENSE +0 -0
- {code_airlock-1.0.0 → code_airlock-1.3.0}/setup.cfg +0 -0
- {code_airlock-1.0.0 → code_airlock-1.3.0}/src/code_airlock.egg-info/dependency_links.txt +0 -0
- {code_airlock-1.0.0 → code_airlock-1.3.0}/src/code_airlock.egg-info/entry_points.txt +0 -0
- {code_airlock-1.0.0 → code_airlock-1.3.0}/src/code_airlock.egg-info/requires.txt +0 -0
- {code_airlock-1.0.0 → code_airlock-1.3.0}/src/code_airlock.egg-info/top_level.txt +0 -0
- {code_airlock-1.0.0 → code_airlock-1.3.0}/tests/test_proxy.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: code-airlock
|
|
3
|
-
Version: 1.
|
|
3
|
+
Version: 1.3.0
|
|
4
4
|
Summary: Low-latency OpenAI-compatible PII redaction gateway and FastMCP server for coding agents
|
|
5
5
|
Author: Sarthak Maurya
|
|
6
6
|
License-Expression: MIT
|
|
@@ -45,8 +45,12 @@ Dynamic: license-file
|
|
|
45
45
|
|
|
46
46
|
* **Dual Deployment Modes:** Run as a drop-in **OpenAI-compatible `/v1/chat/completions` reverse proxy** or as a **Model Context Protocol (MCP) server** (`stdio`) for OpenCode, Cursor, and Claude Code.
|
|
47
47
|
* **2-Stage Hybrid Cascade:** Combines sub-millisecond deterministic regex + mathematical checksums (**Luhn, Verhoeff $D_5$, ISO 7064 MOD-97**) with contextual neural NER (`urchade/gliner_small-v2.1`) and a curated developer allowlist.
|
|
48
|
+
* **Deterministic Secret Coverage:** Stage 1 also captures inline connection-string passwords and `PASSWORD` / `SECRET` / `ACCESS_TOKEN` / `API_KEY` assignments, excluding placeholders, booleans, function calls, and existing vault tokens.
|
|
49
|
+
* **Git & CI Guardrails:** `code-airlock scan --staged` catches secrets in staged files at sub-millisecond Stage 1 speed, masks every value in output, and installs a `.git/hooks/pre-commit` gate via `init --pre-commit`.
|
|
50
|
+
* **Safe Agent Execution:** `airlock_exec` injects real `.env` values into a guarded subprocess (`shell=False`) and scrubs both stdout and stderr before returning them to the model.
|
|
51
|
+
* **Project-Aware Rules:** `.airlock.toml` adds team allowlist terms and custom Stage-1 regexes such as `EMPLOYEE_ID = 'EMP-[0-9]{6}'`.
|
|
48
52
|
* **Zero-Delay Streaming Rehydration:** Stateful sliding-window SSE parser reassembles split tokens across arbitrary 2-byte network chunks without delaying code comparison or bitwise operators (`<`, `<<`, `template <typename T>`).
|
|
49
|
-
* **Minimalist Systems Footprint:** ~
|
|
53
|
+
* **Minimalist Systems Footprint:** ~800 lines of core runtime code (`engine.py`, `sanitizer.py`, `proxy.py`, `mcp_server.py`, `scanner.py`, `config.py`, `executor.py`) plus a deterministic 2,000-row benchmark harness.
|
|
50
54
|
|
|
51
55
|
---
|
|
52
56
|
|
|
@@ -126,7 +130,7 @@ Evaluated on **285 developer prompts** (sampled from `CodeAlpaca_20K` and `OpenC
|
|
|
126
130
|
| `adversarial_sse` | 250 | 250 | 0 / 0 | 1,217 |
|
|
127
131
|
| **Total** | **2,000** | **1,400** | **350 / 250** | **5,817** |
|
|
128
132
|
|
|
129
|
-
Reference result (strict exact-span matching, `code-airlock`
|
|
133
|
+
Reference result (strict exact-span matching, `code-airlock` v1.3.0):
|
|
130
134
|
|
|
131
135
|
| Split | Strict P / R / F1 | Hard-neg | Clean 0-FP | SSE 1/2/7-byte |
|
|
132
136
|
| :--- | :---: | :---: | :---: | :---: |
|
|
@@ -136,7 +140,7 @@ Reference result (strict exact-span matching, `code-airlock` v0.1.0):
|
|
|
136
140
|
| `adversarial_sse` | 100% / 100% / 100% | — | — | 250/250 |
|
|
137
141
|
| **Combined** | **100% / 100% / 100%** | **250/250** | **350/350** | **2000/2000** |
|
|
138
142
|
|
|
139
|
-
- Latency: **0.
|
|
143
|
+
- Latency: **0.232 ms** p50 / 0.428 ms p95 Stage 1; **32.354 ms** p50 / 43.966 ms p95 Stage 2; **32.704 ms** p50 / 44.112 ms p95 total cascade.
|
|
140
144
|
- Per-entity support: `API_KEY` 756, `AUTH_TOKEN` 467, `CREDIT_CARD` 440, `EMAIL` 754, `PHONE` 465, `IP_ADDRESS` 442, `GOVT_ID` 465, `IBAN_CRYPTO` 465, `PERSON` 729, `ORGANIZATION` 417, `LOCATION` 417. Every entity is at **100% P/R/F1**.
|
|
141
145
|
- Invariants: MCP JSON parse **100/100**, JSON round-trip **100/100**, adversarial SSE **2,000/2,000**, hard negatives **250/250**, clean rows **350/350 with zero false positives**.
|
|
142
146
|
|
|
@@ -242,6 +246,77 @@ python -m code_airlock
|
|
|
242
246
|
code-airlock-mcp
|
|
243
247
|
```
|
|
244
248
|
|
|
249
|
+
### Zero-Friction CLI
|
|
250
|
+
|
|
251
|
+
```bash
|
|
252
|
+
python -m code_airlock doctor # model warmup + round-trip + MCP config
|
|
253
|
+
python -m code_airlock init # add rules to .cursorrules / CLAUDE.md
|
|
254
|
+
python -m code_airlock init --pre-commit # also install .git/hooks/pre-commit
|
|
255
|
+
python -m code_airlock init --config # create a starter .airlock.toml
|
|
256
|
+
python -m code_airlock scan --staged # fast Stage 1 scan of staged files
|
|
257
|
+
python -m code_airlock scan --deep src/ # add GLiNER detection for a path
|
|
258
|
+
python -m code_airlock proxy # start the HTTP proxy
|
|
259
|
+
python -m code_airlock # start the MCP server (default)
|
|
260
|
+
```
|
|
261
|
+
|
|
262
|
+
Installed entrypoints: `code-airlock-mcp`, `code-airlock-proxy`, and the
|
|
263
|
+
backward-compatible `code-airlock` alias.
|
|
264
|
+
|
|
265
|
+
`scan` skips binaries, files over 1 MB, `tests/`, and benchmark fixtures. It
|
|
266
|
+
prints `file:line -> [CATEGORY] masked-preview` and exits `1` when secrets are
|
|
267
|
+
found, so it works unchanged in CI and Git hooks. `--deep` enables Stage 2 NER;
|
|
268
|
+
the default Stage 1 path is sub-millisecond. `init` is idempotent, and
|
|
269
|
+
`init --pre-commit` installs a hook that runs `python -m code_airlock scan --staged`.
|
|
270
|
+
|
|
271
|
+
### 🧰 MCP Tools
|
|
272
|
+
|
|
273
|
+
| Tool | Purpose |
|
|
274
|
+
| :--- | :--- |
|
|
275
|
+
| `airlock_read_file` | Read a bounded UTF-8 file through the vault; blocks `.ssh`, `.aws`, `.gnupg`, `.kube`, `/etc`, `System32`, `AIRLOCK_WORKSPACE` escapes, binaries, and files over 1 MB. |
|
|
276
|
+
| `airlock_write_file` | Rehydrate locally and write; refuses to write if unresolved/hallucinated `<ENTITY_N>` tokens remain. |
|
|
277
|
+
| `airlock_scrub_text` | Redact arbitrary in-memory text or terminal output; `use_ner=False` forces the fast Stage 1 path. |
|
|
278
|
+
| `airlock_exec` | Run a guarded subprocess (`shell=False`), inject real `.env` values into its environment, and scrub stdout/stderr before returning. Blocks destructive executables and shells. |
|
|
279
|
+
| `airlock_vault_status` | Return active token keys and per-category counts only; never raw values. |
|
|
280
|
+
| `airlock_dehydrate` | Redact text and return session-stable tokens. |
|
|
281
|
+
| `airlock_rehydrate` | Restore canonical, case-mangled, or whitespace-padded vault tokens. |
|
|
282
|
+
|
|
283
|
+
All file tools resolve paths with `Path.expanduser().resolve()` and jail access
|
|
284
|
+
to `AIRLOCK_WORKSPACE` when that variable is set. Rehydration is tolerant of
|
|
285
|
+
`<api_key_1>` / `< API_KEY_1 >`, while the SSE `StreamRehydrator` remains
|
|
286
|
+
zero-delay and exact.
|
|
287
|
+
|
|
288
|
+
`airlock_exec` is designed for safe test runs:
|
|
289
|
+
|
|
290
|
+
```python
|
|
291
|
+
airlock_exec(
|
|
292
|
+
[sys.executable, "-c", "import os; print(os.environ['STRIPE_SECRET_KEY'])"],
|
|
293
|
+
".env",
|
|
294
|
+
)
|
|
295
|
+
# stdout: "<API_KEY_1>" ← the real key never reaches the model
|
|
296
|
+
```
|
|
297
|
+
|
|
298
|
+
### ⚙️ Project Config (`.airlock.toml`)
|
|
299
|
+
|
|
300
|
+
Create a starter with `python -m code_airlock init --config`, or write one
|
|
301
|
+
manually. `CascadeEngine` loads `.airlock.toml` from the current directory, or
|
|
302
|
+
from `AIRLOCK_CONFIG` when set:
|
|
303
|
+
|
|
304
|
+
```toml
|
|
305
|
+
[allowlist]
|
|
306
|
+
terms = ["MyInternalApp", "staging.corp.local"]
|
|
307
|
+
|
|
308
|
+
[custom_patterns]
|
|
309
|
+
EMPLOYEE_ID = 'EMP-[0-9]{6}'
|
|
310
|
+
INTERNAL_TICKET = 'JIRA-[0-9]{4,6}'
|
|
311
|
+
```
|
|
312
|
+
|
|
313
|
+
- Allowlist terms are merged into the Stage 2 allowlist so they are not
|
|
314
|
+
misclassified as `PERSON`, `ORGANIZATION`, or `LOCATION`.
|
|
315
|
+
- Custom patterns run in Stage 1 alongside built-in rules and produce
|
|
316
|
+
`<EMPLOYEE_ID_n>` / `<INTERNAL_TICKET_n>` vault tokens.
|
|
317
|
+
- Invalid regexes and malformed labels are ignored rather than crashing the
|
|
318
|
+
engine.
|
|
319
|
+
|
|
245
320
|
---
|
|
246
321
|
|
|
247
322
|
### Option 3: OpenAI-Compatible Reverse Proxy (`/v1/chat/completions`)
|
|
@@ -277,15 +352,18 @@ To ensure your coding agent actually leverages the airlock tools instead of raw
|
|
|
277
352
|
|
|
278
353
|
| File | Lines | Responsibility |
|
|
279
354
|
| :--- | :---: | :--- |
|
|
280
|
-
| `src/code_airlock/engine.py` |
|
|
281
|
-
| `src/code_airlock/sanitizer.py` |
|
|
355
|
+
| `src/code_airlock/engine.py` | 219 | Stage 1 regex/checksum validators + high-entropy URI/KV secrets + lazy Stage 2 GLiNER |
|
|
356
|
+
| `src/code_airlock/sanitizer.py` | 108 | Vault hygiene, tolerant rehydration, unresolved-token detection, SSE parser |
|
|
282
357
|
| `src/code_airlock/proxy.py` | 153 | FastAPI `/v1/chat/completions` gateway, SSE rehydration, telemetry headers, `main()` |
|
|
283
|
-
| `src/code_airlock/mcp_server.py` |
|
|
358
|
+
| `src/code_airlock/mcp_server.py` | 124 | FastMCP server with path jailing, file guardrails, `airlock_scrub_text`, `airlock_exec`, and write verification |
|
|
359
|
+
| `src/code_airlock/config.py` | 39 | Optional `.airlock.toml` allowlist and custom Stage-1 pattern loader |
|
|
360
|
+
| `src/code_airlock/executor.py` | 87 | Guarded subprocess execution, env injection, and output scrubbing |
|
|
361
|
+
| `src/code_airlock/scanner.py` | 76 | Masked Stage-1/Stage-2 Git and path secret scanner |
|
|
284
362
|
| `src/code_airlock/__init__.py` | 7 | Public API: `CascadeEngine`, `VaultSession`, `StreamRehydrator`, `__version__` |
|
|
285
|
-
| `src/code_airlock/__main__.py` |
|
|
363
|
+
| `src/code_airlock/__main__.py` | 127 | CLI: `mcp`, `proxy`, `doctor`, `init`, `scan` |
|
|
286
364
|
| `tests/test_proxy.py` | 84 | End-to-end proxy tests (11-entity zero-leak check, 2-char SSE split, C++ syntax) |
|
|
287
|
-
| `tests/test_mcp.py` |
|
|
288
|
-
| `benchmarks/benchmark.py` |
|
|
365
|
+
| `tests/test_mcp.py` | 181 | MCP tools, scanner, exec, config, pre-commit hook, file jail, tolerant rehydration |
|
|
366
|
+
| `benchmarks/benchmark.py` | 163 | Unified 4-suite evaluator → `benchmarks/BENCHMARK_RESULTS.md` |
|
|
289
367
|
| `benchmarks/build_benchmarks.py` | 239 | Deterministic 2,000-row generator (`seed=42`) |
|
|
290
368
|
| `benchmarks/export_hf.py` | 88 | Publishes all 4 splits plus the auto-generated dataset card |
|
|
291
369
|
| `benchmarks/BENCHMARK_RESULTS.md` | — | Generated metrics-only scorecard |
|
|
@@ -303,6 +381,15 @@ Detection is best-effort and local-only. It reduces exposure but is not a
|
|
|
303
381
|
guarantee that every secret is removed. See [SECURITY.md](SECURITY.md) for the
|
|
304
382
|
threat model, known gaps, and private disclosure instructions.
|
|
305
383
|
|
|
384
|
+
The MCP file tools additionally block sensitive directories (`~/.ssh`, `~/.aws`,
|
|
385
|
+
`~/.gnupg`, `~/.kube`, `/etc`, `System32`), files over 1 MB, binary files, and
|
|
386
|
+
paths outside `AIRLOCK_WORKSPACE` when that variable is set. Writes containing
|
|
387
|
+
unresolved or hallucinated `<ENTITY_N>` tokens are refused.
|
|
388
|
+
|
|
389
|
+
`code-airlock scan --staged` adds a Git pre-commit gate for the five highest-risk
|
|
390
|
+
secret categories (`API_KEY`, `AUTH_TOKEN`, `CREDIT_CARD`, `GOVT_ID`,
|
|
391
|
+
`IBAN_CRYPTO`). It prints masked previews only and exits non-zero on findings.
|
|
392
|
+
|
|
306
393
|
## 📜 License
|
|
307
394
|
|
|
308
395
|
MIT © 2026 Sarthak Maurya. See [LICENSE](LICENSE).
|
|
@@ -6,8 +6,12 @@
|
|
|
6
6
|
|
|
7
7
|
* **Dual Deployment Modes:** Run as a drop-in **OpenAI-compatible `/v1/chat/completions` reverse proxy** or as a **Model Context Protocol (MCP) server** (`stdio`) for OpenCode, Cursor, and Claude Code.
|
|
8
8
|
* **2-Stage Hybrid Cascade:** Combines sub-millisecond deterministic regex + mathematical checksums (**Luhn, Verhoeff $D_5$, ISO 7064 MOD-97**) with contextual neural NER (`urchade/gliner_small-v2.1`) and a curated developer allowlist.
|
|
9
|
+
* **Deterministic Secret Coverage:** Stage 1 also captures inline connection-string passwords and `PASSWORD` / `SECRET` / `ACCESS_TOKEN` / `API_KEY` assignments, excluding placeholders, booleans, function calls, and existing vault tokens.
|
|
10
|
+
* **Git & CI Guardrails:** `code-airlock scan --staged` catches secrets in staged files at sub-millisecond Stage 1 speed, masks every value in output, and installs a `.git/hooks/pre-commit` gate via `init --pre-commit`.
|
|
11
|
+
* **Safe Agent Execution:** `airlock_exec` injects real `.env` values into a guarded subprocess (`shell=False`) and scrubs both stdout and stderr before returning them to the model.
|
|
12
|
+
* **Project-Aware Rules:** `.airlock.toml` adds team allowlist terms and custom Stage-1 regexes such as `EMPLOYEE_ID = 'EMP-[0-9]{6}'`.
|
|
9
13
|
* **Zero-Delay Streaming Rehydration:** Stateful sliding-window SSE parser reassembles split tokens across arbitrary 2-byte network chunks without delaying code comparison or bitwise operators (`<`, `<<`, `template <typename T>`).
|
|
10
|
-
* **Minimalist Systems Footprint:** ~
|
|
14
|
+
* **Minimalist Systems Footprint:** ~800 lines of core runtime code (`engine.py`, `sanitizer.py`, `proxy.py`, `mcp_server.py`, `scanner.py`, `config.py`, `executor.py`) plus a deterministic 2,000-row benchmark harness.
|
|
11
15
|
|
|
12
16
|
---
|
|
13
17
|
|
|
@@ -87,7 +91,7 @@ Evaluated on **285 developer prompts** (sampled from `CodeAlpaca_20K` and `OpenC
|
|
|
87
91
|
| `adversarial_sse` | 250 | 250 | 0 / 0 | 1,217 |
|
|
88
92
|
| **Total** | **2,000** | **1,400** | **350 / 250** | **5,817** |
|
|
89
93
|
|
|
90
|
-
Reference result (strict exact-span matching, `code-airlock`
|
|
94
|
+
Reference result (strict exact-span matching, `code-airlock` v1.3.0):
|
|
91
95
|
|
|
92
96
|
| Split | Strict P / R / F1 | Hard-neg | Clean 0-FP | SSE 1/2/7-byte |
|
|
93
97
|
| :--- | :---: | :---: | :---: | :---: |
|
|
@@ -97,7 +101,7 @@ Reference result (strict exact-span matching, `code-airlock` v0.1.0):
|
|
|
97
101
|
| `adversarial_sse` | 100% / 100% / 100% | — | — | 250/250 |
|
|
98
102
|
| **Combined** | **100% / 100% / 100%** | **250/250** | **350/350** | **2000/2000** |
|
|
99
103
|
|
|
100
|
-
- Latency: **0.
|
|
104
|
+
- Latency: **0.232 ms** p50 / 0.428 ms p95 Stage 1; **32.354 ms** p50 / 43.966 ms p95 Stage 2; **32.704 ms** p50 / 44.112 ms p95 total cascade.
|
|
101
105
|
- Per-entity support: `API_KEY` 756, `AUTH_TOKEN` 467, `CREDIT_CARD` 440, `EMAIL` 754, `PHONE` 465, `IP_ADDRESS` 442, `GOVT_ID` 465, `IBAN_CRYPTO` 465, `PERSON` 729, `ORGANIZATION` 417, `LOCATION` 417. Every entity is at **100% P/R/F1**.
|
|
102
106
|
- Invariants: MCP JSON parse **100/100**, JSON round-trip **100/100**, adversarial SSE **2,000/2,000**, hard negatives **250/250**, clean rows **350/350 with zero false positives**.
|
|
103
107
|
|
|
@@ -203,6 +207,77 @@ python -m code_airlock
|
|
|
203
207
|
code-airlock-mcp
|
|
204
208
|
```
|
|
205
209
|
|
|
210
|
+
### Zero-Friction CLI
|
|
211
|
+
|
|
212
|
+
```bash
|
|
213
|
+
python -m code_airlock doctor # model warmup + round-trip + MCP config
|
|
214
|
+
python -m code_airlock init # add rules to .cursorrules / CLAUDE.md
|
|
215
|
+
python -m code_airlock init --pre-commit # also install .git/hooks/pre-commit
|
|
216
|
+
python -m code_airlock init --config # create a starter .airlock.toml
|
|
217
|
+
python -m code_airlock scan --staged # fast Stage 1 scan of staged files
|
|
218
|
+
python -m code_airlock scan --deep src/ # add GLiNER detection for a path
|
|
219
|
+
python -m code_airlock proxy # start the HTTP proxy
|
|
220
|
+
python -m code_airlock # start the MCP server (default)
|
|
221
|
+
```
|
|
222
|
+
|
|
223
|
+
Installed entrypoints: `code-airlock-mcp`, `code-airlock-proxy`, and the
|
|
224
|
+
backward-compatible `code-airlock` alias.
|
|
225
|
+
|
|
226
|
+
`scan` skips binaries, files over 1 MB, `tests/`, and benchmark fixtures. It
|
|
227
|
+
prints `file:line -> [CATEGORY] masked-preview` and exits `1` when secrets are
|
|
228
|
+
found, so it works unchanged in CI and Git hooks. `--deep` enables Stage 2 NER;
|
|
229
|
+
the default Stage 1 path is sub-millisecond. `init` is idempotent, and
|
|
230
|
+
`init --pre-commit` installs a hook that runs `python -m code_airlock scan --staged`.
|
|
231
|
+
|
|
232
|
+
### 🧰 MCP Tools
|
|
233
|
+
|
|
234
|
+
| Tool | Purpose |
|
|
235
|
+
| :--- | :--- |
|
|
236
|
+
| `airlock_read_file` | Read a bounded UTF-8 file through the vault; blocks `.ssh`, `.aws`, `.gnupg`, `.kube`, `/etc`, `System32`, `AIRLOCK_WORKSPACE` escapes, binaries, and files over 1 MB. |
|
|
237
|
+
| `airlock_write_file` | Rehydrate locally and write; refuses to write if unresolved/hallucinated `<ENTITY_N>` tokens remain. |
|
|
238
|
+
| `airlock_scrub_text` | Redact arbitrary in-memory text or terminal output; `use_ner=False` forces the fast Stage 1 path. |
|
|
239
|
+
| `airlock_exec` | Run a guarded subprocess (`shell=False`), inject real `.env` values into its environment, and scrub stdout/stderr before returning. Blocks destructive executables and shells. |
|
|
240
|
+
| `airlock_vault_status` | Return active token keys and per-category counts only; never raw values. |
|
|
241
|
+
| `airlock_dehydrate` | Redact text and return session-stable tokens. |
|
|
242
|
+
| `airlock_rehydrate` | Restore canonical, case-mangled, or whitespace-padded vault tokens. |
|
|
243
|
+
|
|
244
|
+
All file tools resolve paths with `Path.expanduser().resolve()` and jail access
|
|
245
|
+
to `AIRLOCK_WORKSPACE` when that variable is set. Rehydration is tolerant of
|
|
246
|
+
`<api_key_1>` / `< API_KEY_1 >`, while the SSE `StreamRehydrator` remains
|
|
247
|
+
zero-delay and exact.
|
|
248
|
+
|
|
249
|
+
`airlock_exec` is designed for safe test runs:
|
|
250
|
+
|
|
251
|
+
```python
|
|
252
|
+
airlock_exec(
|
|
253
|
+
[sys.executable, "-c", "import os; print(os.environ['STRIPE_SECRET_KEY'])"],
|
|
254
|
+
".env",
|
|
255
|
+
)
|
|
256
|
+
# stdout: "<API_KEY_1>" ← the real key never reaches the model
|
|
257
|
+
```
|
|
258
|
+
|
|
259
|
+
### ⚙️ Project Config (`.airlock.toml`)
|
|
260
|
+
|
|
261
|
+
Create a starter with `python -m code_airlock init --config`, or write one
|
|
262
|
+
manually. `CascadeEngine` loads `.airlock.toml` from the current directory, or
|
|
263
|
+
from `AIRLOCK_CONFIG` when set:
|
|
264
|
+
|
|
265
|
+
```toml
|
|
266
|
+
[allowlist]
|
|
267
|
+
terms = ["MyInternalApp", "staging.corp.local"]
|
|
268
|
+
|
|
269
|
+
[custom_patterns]
|
|
270
|
+
EMPLOYEE_ID = 'EMP-[0-9]{6}'
|
|
271
|
+
INTERNAL_TICKET = 'JIRA-[0-9]{4,6}'
|
|
272
|
+
```
|
|
273
|
+
|
|
274
|
+
- Allowlist terms are merged into the Stage 2 allowlist so they are not
|
|
275
|
+
misclassified as `PERSON`, `ORGANIZATION`, or `LOCATION`.
|
|
276
|
+
- Custom patterns run in Stage 1 alongside built-in rules and produce
|
|
277
|
+
`<EMPLOYEE_ID_n>` / `<INTERNAL_TICKET_n>` vault tokens.
|
|
278
|
+
- Invalid regexes and malformed labels are ignored rather than crashing the
|
|
279
|
+
engine.
|
|
280
|
+
|
|
206
281
|
---
|
|
207
282
|
|
|
208
283
|
### Option 3: OpenAI-Compatible Reverse Proxy (`/v1/chat/completions`)
|
|
@@ -238,15 +313,18 @@ To ensure your coding agent actually leverages the airlock tools instead of raw
|
|
|
238
313
|
|
|
239
314
|
| File | Lines | Responsibility |
|
|
240
315
|
| :--- | :---: | :--- |
|
|
241
|
-
| `src/code_airlock/engine.py` |
|
|
242
|
-
| `src/code_airlock/sanitizer.py` |
|
|
316
|
+
| `src/code_airlock/engine.py` | 219 | Stage 1 regex/checksum validators + high-entropy URI/KV secrets + lazy Stage 2 GLiNER |
|
|
317
|
+
| `src/code_airlock/sanitizer.py` | 108 | Vault hygiene, tolerant rehydration, unresolved-token detection, SSE parser |
|
|
243
318
|
| `src/code_airlock/proxy.py` | 153 | FastAPI `/v1/chat/completions` gateway, SSE rehydration, telemetry headers, `main()` |
|
|
244
|
-
| `src/code_airlock/mcp_server.py` |
|
|
319
|
+
| `src/code_airlock/mcp_server.py` | 124 | FastMCP server with path jailing, file guardrails, `airlock_scrub_text`, `airlock_exec`, and write verification |
|
|
320
|
+
| `src/code_airlock/config.py` | 39 | Optional `.airlock.toml` allowlist and custom Stage-1 pattern loader |
|
|
321
|
+
| `src/code_airlock/executor.py` | 87 | Guarded subprocess execution, env injection, and output scrubbing |
|
|
322
|
+
| `src/code_airlock/scanner.py` | 76 | Masked Stage-1/Stage-2 Git and path secret scanner |
|
|
245
323
|
| `src/code_airlock/__init__.py` | 7 | Public API: `CascadeEngine`, `VaultSession`, `StreamRehydrator`, `__version__` |
|
|
246
|
-
| `src/code_airlock/__main__.py` |
|
|
324
|
+
| `src/code_airlock/__main__.py` | 127 | CLI: `mcp`, `proxy`, `doctor`, `init`, `scan` |
|
|
247
325
|
| `tests/test_proxy.py` | 84 | End-to-end proxy tests (11-entity zero-leak check, 2-char SSE split, C++ syntax) |
|
|
248
|
-
| `tests/test_mcp.py` |
|
|
249
|
-
| `benchmarks/benchmark.py` |
|
|
326
|
+
| `tests/test_mcp.py` | 181 | MCP tools, scanner, exec, config, pre-commit hook, file jail, tolerant rehydration |
|
|
327
|
+
| `benchmarks/benchmark.py` | 163 | Unified 4-suite evaluator → `benchmarks/BENCHMARK_RESULTS.md` |
|
|
250
328
|
| `benchmarks/build_benchmarks.py` | 239 | Deterministic 2,000-row generator (`seed=42`) |
|
|
251
329
|
| `benchmarks/export_hf.py` | 88 | Publishes all 4 splits plus the auto-generated dataset card |
|
|
252
330
|
| `benchmarks/BENCHMARK_RESULTS.md` | — | Generated metrics-only scorecard |
|
|
@@ -264,6 +342,15 @@ Detection is best-effort and local-only. It reduces exposure but is not a
|
|
|
264
342
|
guarantee that every secret is removed. See [SECURITY.md](SECURITY.md) for the
|
|
265
343
|
threat model, known gaps, and private disclosure instructions.
|
|
266
344
|
|
|
345
|
+
The MCP file tools additionally block sensitive directories (`~/.ssh`, `~/.aws`,
|
|
346
|
+
`~/.gnupg`, `~/.kube`, `/etc`, `System32`), files over 1 MB, binary files, and
|
|
347
|
+
paths outside `AIRLOCK_WORKSPACE` when that variable is set. Writes containing
|
|
348
|
+
unresolved or hallucinated `<ENTITY_N>` tokens are refused.
|
|
349
|
+
|
|
350
|
+
`code-airlock scan --staged` adds a Git pre-commit gate for the five highest-risk
|
|
351
|
+
secret categories (`API_KEY`, `AUTH_TOKEN`, `CREDIT_CARD`, `GOVT_ID`,
|
|
352
|
+
`IBAN_CRYPTO`). It prints masked previews only and exits non-zero on findings.
|
|
353
|
+
|
|
267
354
|
## 📜 License
|
|
268
355
|
|
|
269
356
|
MIT © 2026 Sarthak Maurya. See [LICENSE](LICENSE).
|
|
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "code-airlock"
|
|
7
|
-
version = "1.
|
|
7
|
+
version = "1.3.0"
|
|
8
8
|
description = "Low-latency OpenAI-compatible PII redaction gateway and FastMCP server for coding agents"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.10"
|
|
@@ -0,0 +1,127 @@
|
|
|
1
|
+
"""Zero-friction CLI: MCP server, proxy, doctor, init, and secret scanning."""
|
|
2
|
+
from __future__ import annotations
|
|
3
|
+
import argparse, os, sys
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
|
|
6
|
+
RULE = """# code-airlock security rules
|
|
7
|
+
Always use `airlock_read_file` and `airlock_write_file` when reading or writing `.env`, `*secret*`, credentials, or deployment configs. Never inspect secret files using raw bash (`cat`, `grep`) or native file read tools.
|
|
8
|
+
"""
|
|
9
|
+
MARKER = "code-airlock security rules"
|
|
10
|
+
HOOK = """#!/bin/sh
|
|
11
|
+
# Installed by code-airlock
|
|
12
|
+
python -m code_airlock scan --staged
|
|
13
|
+
"""
|
|
14
|
+
CONFIG_TEMPLATE = """# code-airlock project configuration
|
|
15
|
+
[allowlist]
|
|
16
|
+
terms = ["MyInternalApp", "staging.corp.local"]
|
|
17
|
+
|
|
18
|
+
[custom_patterns]
|
|
19
|
+
# EMPLOYEE_ID = 'EMP-[0-9]{6}'
|
|
20
|
+
# INTERNAL_TICKET = 'JIRA-[0-9]{4,6}'
|
|
21
|
+
"""
|
|
22
|
+
CONFIG = """{
|
|
23
|
+
"mcpServers": {
|
|
24
|
+
"code-airlock": {
|
|
25
|
+
"command": "python",
|
|
26
|
+
"args": ["-m", "code_airlock.mcp_server"]
|
|
27
|
+
}
|
|
28
|
+
}
|
|
29
|
+
}"""
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
def doctor() -> int:
|
|
33
|
+
from .engine import CascadeEngine
|
|
34
|
+
from .sanitizer import VaultSession
|
|
35
|
+
sample = ("API key sk-" + "proj-abc123def456ghi789jkl012mno345; "
|
|
36
|
+
"postgres://user:" + "pass12345" + "@10.0.0.5:5432/db; "
|
|
37
|
+
"contact alice@acme.io and Alice Smith.")
|
|
38
|
+
engine = CascadeEngine(lazy=False)
|
|
39
|
+
spans = engine.detect(sample)
|
|
40
|
+
labels = {s.label for s in spans}
|
|
41
|
+
vault = VaultSession()
|
|
42
|
+
dehydrated, _ = vault.dehydrate(sample, spans)
|
|
43
|
+
restored = vault.rehydrate(dehydrated)
|
|
44
|
+
expected = {"API_KEY", "AUTH_TOKEN", "EMAIL", "PERSON"}
|
|
45
|
+
missing = expected - labels
|
|
46
|
+
ok = restored == sample and not missing
|
|
47
|
+
print(f"[{'ok' if not missing else 'fail'}] GLiNER ready on {engine.device}; labels: {', '.join(sorted(labels))}")
|
|
48
|
+
print(f"[{'ok' if restored == sample else 'fail'}] round-trip rehydration ({len(spans)} spans)")
|
|
49
|
+
if not ok:
|
|
50
|
+
if missing: print(f"[fail] missing detectors: {', '.join(sorted(missing))}")
|
|
51
|
+
return 1
|
|
52
|
+
print("\nOpenCode (~/.config/opencode/opencode.json):")
|
|
53
|
+
print('{"mcp": {"code-airlock": {"type": "local", "command": ["python", "-m", "code_airlock.mcp_server"]}}}')
|
|
54
|
+
print("\nCursor / Claude Code / Claude Desktop:")
|
|
55
|
+
print(CONFIG)
|
|
56
|
+
return 0
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
def init(pre_commit: bool = False, config: bool = False) -> int:
|
|
60
|
+
if config:
|
|
61
|
+
config_path = Path.cwd() / ".airlock.toml"
|
|
62
|
+
if config_path.exists():
|
|
63
|
+
print("[ok] .airlock.toml already exists")
|
|
64
|
+
else:
|
|
65
|
+
config_path.write_text(CONFIG_TEMPLATE, encoding="utf-8")
|
|
66
|
+
print("[ok] created .airlock.toml")
|
|
67
|
+
for name in (".cursorrules", "CLAUDE.md"):
|
|
68
|
+
target = Path.cwd() / name
|
|
69
|
+
existing = target.read_text(encoding="utf-8") if target.exists() else ""
|
|
70
|
+
if MARKER in existing:
|
|
71
|
+
print(f"[ok] {name} already contains code-airlock rules")
|
|
72
|
+
continue
|
|
73
|
+
separator = "\n\n" if existing.strip() else ""
|
|
74
|
+
target.write_text(existing.rstrip() + separator + RULE, encoding="utf-8")
|
|
75
|
+
print(f"[ok] updated {name}")
|
|
76
|
+
if pre_commit:
|
|
77
|
+
hooks = Path.cwd() / ".git" / "hooks"
|
|
78
|
+
if not hooks.is_dir():
|
|
79
|
+
print("[skip] no .git/hooks directory")
|
|
80
|
+
else:
|
|
81
|
+
hook = hooks / "pre-commit"
|
|
82
|
+
existing = hook.read_text(encoding="utf-8") if hook.exists() else ""
|
|
83
|
+
if "code-airlock scan --staged" in existing:
|
|
84
|
+
print("[ok] .git/hooks/pre-commit already installed")
|
|
85
|
+
else:
|
|
86
|
+
separator = "\n\n" if existing.strip() else ""
|
|
87
|
+
hook.write_text(existing.rstrip() + separator + HOOK, encoding="utf-8")
|
|
88
|
+
os.chmod(hook, 0o755)
|
|
89
|
+
print("[ok] installed .git/hooks/pre-commit")
|
|
90
|
+
return 0
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
def main() -> int:
|
|
94
|
+
parser = argparse.ArgumentParser(prog="code-airlock", description="Local PII/secret gate for coding agents")
|
|
95
|
+
sub = parser.add_subparsers(dest="command")
|
|
96
|
+
sub.add_parser("mcp", help="start the MCP stdio server (default)")
|
|
97
|
+
sub.add_parser("proxy", help="start the OpenAI-compatible HTTP proxy")
|
|
98
|
+
sub.add_parser("doctor", help="verify the model and print MCP configuration")
|
|
99
|
+
init_cmd = sub.add_parser("init", help="add code-airlock rules and optional pre-commit hook")
|
|
100
|
+
init_cmd.add_argument("--pre-commit", action="store_true", help="install .git/hooks/pre-commit")
|
|
101
|
+
init_cmd.add_argument("--config", action="store_true", help="create a starter .airlock.toml")
|
|
102
|
+
scan_cmd = sub.add_parser("scan", help="scan files or staged changes for secrets")
|
|
103
|
+
scan_cmd.add_argument("--staged", action="store_true", help="scan files staged in git")
|
|
104
|
+
scan_cmd.add_argument("--deep", action="store_true", help="enable Stage 2 GLiNER detection")
|
|
105
|
+
scan_cmd.add_argument("paths", nargs="*", default=["."], help="files or directories to scan")
|
|
106
|
+
args = parser.parse_args()
|
|
107
|
+
if args.command in (None, "mcp"):
|
|
108
|
+
from .mcp_server import main as run
|
|
109
|
+
run()
|
|
110
|
+
return 0
|
|
111
|
+
if args.command == "proxy":
|
|
112
|
+
from .proxy import main as run
|
|
113
|
+
run()
|
|
114
|
+
return 0
|
|
115
|
+
if args.command == "doctor":
|
|
116
|
+
return doctor()
|
|
117
|
+
if args.command == "init":
|
|
118
|
+
return init(args.pre_commit, args.config)
|
|
119
|
+
if args.command == "scan":
|
|
120
|
+
from .scanner import scan
|
|
121
|
+
return scan(args.paths, staged=args.staged, deep=args.deep)
|
|
122
|
+
parser.print_help()
|
|
123
|
+
return 0
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
if __name__ == "__main__":
|
|
127
|
+
sys.exit(main())
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
"""config.py - Optional .airlock.toml allowlist and custom Stage-1 patterns."""
|
|
2
|
+
from __future__ import annotations
|
|
3
|
+
import os, re
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
|
|
6
|
+
try:
|
|
7
|
+
import tomllib
|
|
8
|
+
except ImportError: # Python 3.10 fallback; config is optional
|
|
9
|
+
try:
|
|
10
|
+
import tomli as tomllib
|
|
11
|
+
except ImportError:
|
|
12
|
+
tomllib = None
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def load_config(path: str | Path | None = None):
|
|
16
|
+
"""Return (allowlist terms, compiled custom patterns) from .airlock.toml."""
|
|
17
|
+
if tomllib is None:
|
|
18
|
+
return frozenset(), {}
|
|
19
|
+
target = Path(path or os.getenv("AIRLOCK_CONFIG") or ".airlock.toml")
|
|
20
|
+
if not target.is_file():
|
|
21
|
+
return frozenset(), {}
|
|
22
|
+
try:
|
|
23
|
+
with target.open("rb") as handle:
|
|
24
|
+
data = tomllib.load(handle)
|
|
25
|
+
except (OSError, ValueError):
|
|
26
|
+
return frozenset(), {}
|
|
27
|
+
terms = data.get("allowlist", {})
|
|
28
|
+
terms = terms.get("terms", []) if isinstance(terms, dict) else []
|
|
29
|
+
allowlist = frozenset(str(term).lower() for term in terms) if isinstance(terms, list) else frozenset()
|
|
30
|
+
patterns = {}
|
|
31
|
+
patterns_data = data.get("custom_patterns", {})
|
|
32
|
+
for label, raw in (patterns_data.items() if isinstance(patterns_data, dict) else []):
|
|
33
|
+
if not re.fullmatch(r"[A-Z_]+", label) or not isinstance(raw, str):
|
|
34
|
+
continue
|
|
35
|
+
try:
|
|
36
|
+
patterns[label] = re.compile(raw)
|
|
37
|
+
except re.error:
|
|
38
|
+
continue
|
|
39
|
+
return allowlist, patterns
|
|
@@ -2,8 +2,7 @@
|
|
|
2
2
|
from __future__ import annotations
|
|
3
3
|
import re
|
|
4
4
|
from dataclasses import dataclass
|
|
5
|
-
import
|
|
6
|
-
from gliner import GLiNER
|
|
5
|
+
from .config import load_config
|
|
7
6
|
|
|
8
7
|
@dataclass(frozen=True)
|
|
9
8
|
class Span:
|
|
@@ -77,6 +76,11 @@ PAN_RE = re.compile(r"\b[A-Z]{5}[0-9]{4}[A-Z]\b")
|
|
|
77
76
|
AADHAAR_RE = re.compile(r"\b[2-9]\d{3}[ -]?\d{4}[ -]?\d{4}\b")
|
|
78
77
|
ETH_RE = re.compile(r"\b0x[a-fA-F0-9]{40}\b")
|
|
79
78
|
IBAN_RE = re.compile(r"\b[A-Z]{2}\d{2}[A-Z0-9]{11,30}\b")
|
|
79
|
+
URI_PASSWORD_RE = re.compile(r"://[^:\s/@]+:([^@\s/:]{8,})@")
|
|
80
|
+
KV_SECRET_RE = re.compile(r"(?i)(?:PASSWORD|PASSWD|SECRET_KEY|SECRET|PRIVATE_KEY|ACCESS_TOKEN|AUTH_TOKEN|API_KEY)\s*[=:]\s*[\"']?([^\s\"'\r\n#;,]{8,})[\"']?")
|
|
81
|
+
|
|
82
|
+
def _secret_ok(value: str) -> bool:
|
|
83
|
+
return value.lower() not in {"true", "false", "null", "none", "changeme", "localhost"} and not value.startswith(("<", "$", "{")) and not any(c in value for c in "(){}") and len(set(value)) > 2
|
|
80
84
|
|
|
81
85
|
TECH_ALLOWLIST = frozenset({
|
|
82
86
|
"docker", "redis", "fastapi", "pytorch", "react", "aws", "postgres", "postgresql",
|
|
@@ -102,20 +106,18 @@ PERSON_STOPWORDS = frozenset({
|
|
|
102
106
|
})
|
|
103
107
|
|
|
104
108
|
CODE_CHARS = frozenset("_()[]{}:;/\\=$@")
|
|
105
|
-
SCHEMA_WORDS = frozenset({"production", "development", "staging", "prod", "dev", "test", "localhost",
|
|
106
|
-
"api", "backend", "frontend", "database", "server", "client", "service", "origin",
|
|
107
|
-
"environment", "region", "payments", "config", "metadata"})
|
|
109
|
+
SCHEMA_WORDS = frozenset({"production", "development", "staging", "prod", "dev", "test", "localhost", "api", "backend", "frontend", "database", "server", "client", "service", "origin", "environment", "region", "payments", "config", "metadata"})
|
|
108
110
|
|
|
109
|
-
def stage2_ok(text: str, label: str) -> bool:
|
|
111
|
+
def stage2_ok(text: str, label: str, allowlist: frozenset = TECH_ALLOWLIST) -> bool:
|
|
110
112
|
"""Reject Stage 2 drift on code identifiers, tech names, and role words."""
|
|
111
113
|
t = text.strip()
|
|
112
114
|
if len(t) <= 2: return False
|
|
113
115
|
low = t.lower()
|
|
114
|
-
if low in
|
|
116
|
+
if low in allowlist or low in {"location", "organization", "person", "address", "user"}: return False
|
|
115
117
|
if label in ("LOCATION", "ORGANIZATION") and low in SCHEMA_WORDS: return False
|
|
116
118
|
if label == "LOCATION" and (re.fullmatch(r"[a-z]{2}(?:-[a-z]+)+-\d+", low) or any(w in low.split() for w in ("cluster", "node", "proxy", "channel", "region"))): return False
|
|
117
119
|
if label == "ORGANIZATION" and re.fullmatch(r"[a-z0-9]+(?:-[a-z0-9]+)+", t): return False
|
|
118
|
-
if (label == "ORGANIZATION" and (low.endswith(" service") or low.replace(".", "") in
|
|
120
|
+
if (label == "ORGANIZATION" and (low.endswith(" service") or low.replace(".", "") in allowlist)) or re.fullmatch(r"[A-Z]{2}\d{2}[A-Z0-9]{11,30}", t): return False
|
|
119
121
|
if label == "PERSON" and low in PERSON_STOPWORDS: return False
|
|
120
122
|
if label == "PERSON" and " " not in t and (len(t) > 25 or any(c.isdigit() for c in t) or sum(c.isupper() for c in t) / len(t) > 0.5): return False
|
|
121
123
|
if any(ch in CODE_CHARS for ch in t) or (any(c.isdigit() for c in t) and "." in t) or (label in ("LOCATION", "PERSON", "ORGANIZATION") and t.isdigit()): return False
|
|
@@ -157,18 +159,35 @@ def merge_spans(stage1_spans: list[Span], stage2_spans: list[Span]) -> list[Span
|
|
|
157
159
|
return sorted(s1 + s2, key=lambda s: s.start)
|
|
158
160
|
|
|
159
161
|
class CascadeEngine:
|
|
160
|
-
def __init__(self, device: str | None = None, model_name: str = "urchade/gliner_small-v2.1",
|
|
161
|
-
|
|
162
|
+
def __init__(self, device: str | None = None, model_name: str = "urchade/gliner_small-v2.1",
|
|
163
|
+
threshold: float = 0.5, lazy: bool = True):
|
|
164
|
+
self.device = device
|
|
165
|
+
self.model_name = model_name
|
|
162
166
|
self.threshold = threshold
|
|
163
|
-
self.
|
|
164
|
-
if hasattr(self.gliner, "to"): self.gliner.to(self.device)
|
|
167
|
+
extra, self.custom_patterns = load_config(); self.allowlist = TECH_ALLOWLIST | extra
|
|
165
168
|
self.label_map = {"person": "PERSON", "organization": "ORGANIZATION", "location": "LOCATION"}
|
|
169
|
+
self._ner_model = None
|
|
170
|
+
if not lazy: self._get_ner()
|
|
171
|
+
|
|
172
|
+
def _get_ner(self):
|
|
173
|
+
if self._ner_model is None:
|
|
174
|
+
import torch
|
|
175
|
+
from gliner import GLiNER
|
|
176
|
+
self.device = self.device or ("cuda" if torch.cuda.is_available() else "cpu")
|
|
177
|
+
self._ner_model = GLiNER.from_pretrained(self.model_name)
|
|
178
|
+
if hasattr(self._ner_model, "to"): self._ner_model.to(self.device)
|
|
179
|
+
return self._ner_model
|
|
166
180
|
|
|
167
181
|
def detect_stage1(self, text: str) -> list[Span]:
|
|
168
182
|
spans: list[Span] = []
|
|
169
183
|
for m in EMAIL_RE.finditer(text): spans.append(Span(m.start(), m.end(), "EMAIL", m.group(), 1.0, "stage1"))
|
|
170
184
|
for m in API_KEY_RE.finditer(text): spans.append(Span(m.start(), m.end(), "API_KEY", m.group(), 1.0, "stage1"))
|
|
171
185
|
for m in AUTH_TOKEN_RE.finditer(text): spans.append(Span(m.start(), m.end(), "AUTH_TOKEN", m.group(), 1.0, "stage1"))
|
|
186
|
+
for m in URI_PASSWORD_RE.finditer(text):
|
|
187
|
+
if _secret_ok(m.group(1)): spans.append(Span(m.start(1), m.end(1), "AUTH_TOKEN", m.group(1), 1.0, "stage1"))
|
|
188
|
+
for m in KV_SECRET_RE.finditer(text):
|
|
189
|
+
if _secret_ok(m.group(1)): spans.append(Span(m.start(1), m.end(1), "API_KEY", m.group(1), 1.0, "stage1"))
|
|
190
|
+
for label, pat in self.custom_patterns.items(): spans.extend(Span(m.start(), m.end(), label, m.group(), 1.0, "stage1") for m in pat.finditer(text) if m.group().lower() not in self.allowlist)
|
|
172
191
|
for m in PHONE_RE.finditer(text):
|
|
173
192
|
if is_valid_phone(m.group()): spans.append(Span(m.start(), m.end(), "PHONE", m.group(), 1.0, "stage1"))
|
|
174
193
|
for m in IPV4_CAND_RE.finditer(text):
|
|
@@ -185,15 +204,16 @@ class CascadeEngine:
|
|
|
185
204
|
return _resolve_same_stage(spans)
|
|
186
205
|
|
|
187
206
|
def detect_stage2(self, text: str) -> list[Span]:
|
|
188
|
-
|
|
207
|
+
model = self._get_ner()
|
|
208
|
+
preds = model.predict_entities(text, ["person", "organization", "location"], threshold=self.threshold)
|
|
189
209
|
raw_spans: list[Span] = []
|
|
190
210
|
for p in preds:
|
|
191
211
|
lbl = self.label_map.get(p["label"], p["label"].upper())
|
|
192
|
-
if not stage2_ok(p["text"], lbl): continue
|
|
212
|
+
if not stage2_ok(p["text"], lbl, self.allowlist): continue
|
|
193
213
|
raw_spans.append(Span(p["start"], p["end"], lbl, p["text"], float(p.get("score", 0.0)), "stage2"))
|
|
194
214
|
resolved = _resolve_same_stage(raw_spans)
|
|
195
215
|
resolved_sorted = sorted(resolved, key=lambda s: s.start)
|
|
196
216
|
return merge_adjacent_locations(resolved_sorted, text)
|
|
197
217
|
|
|
198
|
-
def detect(self, text: str) -> list[Span]:
|
|
199
|
-
|
|
218
|
+
def detect(self, text: str, use_ner: bool = True) -> list[Span]:
|
|
219
|
+
stage1 = self.detect_stage1(text); return merge_spans(stage1, self.detect_stage2(text)) if use_ner else stage1
|