# x-3306/onnxstego architecture

> Architecture diagram of the GitHub repository x-3306/onnxstego, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-24). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/x-3306/onnxstego
- Repository: https://github.com/x-3306/onnxstego

## Overview

1. **Repository purpose**
   - This is a Python proof-of-concept for steganographically hiding short authenticated messages inside `float32` ONNX model weights.
   - The repository is organized around a single pipeline: plaintext message → authenticated ciphertext → bit stream → selected ONNX weight LSBs → stego model, and the reverse for extraction.

2. **Primary package boundary: `src/onnxstego/`**
   - The package is split into small, purpose-specific modules rather than one monolith.
   - `core.py` appears to be the orchestration layer that ties together model I/O, bit packing, crypto, and selection.
   - `__main__.py` and `cli.py` provide the command-line entrypoint for the public tool.

3. **Payload protection layer: `crypto.py`**
   - Cryptography is explicitly separated from steganography.
   - The README states the payload uses ChaCha20-Poly1305, so this layer is responsible for authenticated encryption, decryption, and wrong-key/tamper rejection.
   - Architecturally, this means the model only hides ciphertext; message confidentiality and integrity are handled before embedding.

4. **Bit-level encoding layer: `bits.py`**
   - This module likely converts between bytes and bit streams and reconstructs bytes during extraction.
   - It is the bridge between the cryptographic payload format and the per-weight bit edits applied to the ONNX file.

5. **ONNX model handling boundary: `onnx_io.py`**
   - This module is the file-format adapter for reading and writing ONNX tensors.
   - The README highlights support for both `raw_data` and `float_data`, plus normalization to little-endian `float32` before touching mantissa bits.
   - This is a critical boundary because it isolates stego logic from ONNX serialization details and preserves model validity.

6. **Carrier capacity and write locations: `positions.py` and `selection.py`**
   - `positions.py` is responsible for deriving deterministic embedding positions from the master key using a ChaCha20-backed CSPRNG and unbiased sampling.
   - `selection.py` defines which weights are eligible targets.
   - Two selection modes shape the architecture:
     - `uniform`: keyed positions across the full `float32` weight space.
     - `natural`: keyed positions restricted to weights that already differ from a reference model by at least a configured threshold.
   - This is the main “where do bits go?” decision point.

7. **Reference-model dependency for natural selection**
   - The natural mode introduces an external dependency on a reference ONNX model.
   - Sender and receiver must share the same reference identity and threshold, so extraction depends on both the stego model and the reference model.
   - Architecturally, this creates a two-model data flow: reference model + cover model → selection mask → stego embedding.

8. **Top-level workflow in `core.py`**
   - The core flow is likely:
     - inspect model capacity,
     - encrypt message with key and nonce/salt,
     - derive header/payload positions,
     - convert ciphertext to bits,
     - mutate selected mantissa LSBs,
     - write the stego ONNX model.
   - Extraction reverses this flow and stops at AEAD verification if anything is wrong.
   - The README notes that header positions are static per master key while payload positions are re-derived per message from the key plus a fresh salt, so there is a clear split between framing metadata and message content.

9. **CLI as the user-facing boundary**
   - `cli.py` exposes operations such as key generation, inspect, embed, and extract.
   - The CLI is the only intended operational interface; everything else is library plumbing.
   - `__main__.py` makes the package directly runnable as a module.

10. **Reproducibility and proof artifacts**
    - The `proof/` directory contains a complete public demonstration: reference model, cover model, stego model, and a manifest with hashes and command metadata.
    - This is not runtime code; it is an artifact bundle used to prove the approach on a real ONNX model.
    - The proof emphasizes reproducibility, deterministic surrogate fine-tuning, and public verification.

11. **Tooling and empirical validation**
    - `scripts/build_proof.py` reconstructs the public demonstration assets.
    - `scripts/empirical_demo.py` runs a smoke/measurement demo, showing the project is not just a codec but also an analysis tool.
    - The package therefore includes both operational paths and empirical measurement paths.

12. **Test architecture**
    - The `tests/` folder is broad and maps directly to the subsystems:
      - bit packing/unpacking,
      - crypto correctness and rejection behavior,
      - ONNX I/O handling,
      - position derivation,
      - selection logic,
      - end-to-end round trip,
      - empirical/statistical checks.
    - This reflects the repo’s main risk areas: serialization correctness, deterministic keyed selection, and authentication failure modes.

13. **External technology and runtime**
    - The project is Python-based and centered on ONNX model files rather than a custom container format.
    - It relies on ONNX tooling and float32 tensor representation as the carrier medium.
    - The cryptographic and selection design choices are important architecturally because they make the hidden payload keyed, authenticated, and reproducible.

14. **Key architectural takeaway**
    - The repo’s core boundary is between three separable concerns:
      - **message security** in `crypto.py`,
      - **bit/position mapping** in `bits.py`, `positions.py`, and `selection.py`,
      - **model serialization** in `onnx_io.py`.
    - `core.py` composes those pieces into a one-way embed pipeline and a reverse extraction pipeline, while the proof, scripts, and tests validate that the composed system works on real ONNX artifacts.

Main entry points: [`src/onnxstego/core.py`](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/core.py) (Core flow), [`src/onnxstego/selection.py`](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/selection.py) (Selection), [`src/onnxstego/onnx_io.py`](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/onnx_io.py) (ONNX I/O), [`src/onnxstego/positions.py`](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/positions.py) (Positions)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_src_pkg["onnxstego package"]
  node_main["Module entry<br/>entrypoint<br/>[__main__.py]"]
  node_cli["CLI<br/>command interface<br/>[cli.py]"]
  node_core{{"Core flow<br/>orchestrator<br/>[core.py]"}}
  node_crypto["Crypto<br/>aead layer<br/>[crypto.py]"]
  node_bits["Bit codec<br/>bit packing<br/>[bits.py]"]
  node_onnx_io[("ONNX I/O<br/>tensor adapter<br/>[onnx_io.py]")]
  node_positions(("Positions<br/>keyed sampler<br/>[positions.py]"))
  node_selection["Selection<br/>carrier filter<br/>[selection.py]"]
  node_empirical["Empirical<br/>analysis module<br/>[empirical.py]"]
end

subgraph group_tooling["Tooling"]
  node_build_proof["Build proof<br/>demo script<br/>[build_proof.py]"]
  node_demo_script["Demo script<br/>measurement script<br/>[empirical_demo.py]"]
end

subgraph group_artifacts["Proof bundle"]
  node_proof_bundle[("Proof assets<br/>artifact bundle<br/>[README.md]")]
end

subgraph group_validation["Tests"]
  node_test_bits["Bits tests<br/>test suite<br/>[test_bits.py]"]
  node_test_crypto["Crypto tests<br/>test suite<br/>[test_crypto.py]"]
  node_test_onnx["ONNX tests<br/>test suite<br/>[test_onnx_io.py]"]
  node_test_positions["Position tests<br/>test suite<br/>[test_positions.py]"]
  node_test_selection["Selection tests<br/>test suite<br/>[test_selection.py]"]
  node_test_roundtrip{{"Round-trip<br/>end-to-end test"}}
  node_test_empirical["Empirical tests<br/>analysis test<br/>[test_empirical.py]"]
end

node_main -->|"launches"| node_cli
node_cli -->|"invokes"| node_core
node_core -->|"encrypts"| node_crypto
node_core -->|"packs"| node_bits
node_core -->|"reads/writes"| node_onnx_io
node_core -->|"derives"| node_positions
node_core -->|"filters"| node_selection
node_selection -->|"depends on"| node_onnx_io
node_positions -->|"feeds"| node_selection
node_build_proof -->|"produces"| node_proof_bundle
node_demo_script -->|"runs"| node_empirical
node_empirical -->|"measures"| node_core
node_test_bits -.->|"covers"| node_bits
node_test_crypto -.->|"covers"| node_crypto
node_test_onnx -.->|"covers"| node_onnx_io
node_test_positions -.->|"covers"| node_positions
node_test_selection -.->|"covers"| node_selection
node_test_roundtrip -.->|"covers"| node_core
node_test_empirical -.->|"covers"| node_empirical
node_proof_bundle -.->|"demonstrates"| node_onnx_io

click node_main "https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/__main__.py"
click node_cli "https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/cli.py"
click node_core "https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/core.py"
click node_crypto "https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/crypto.py"
click node_bits "https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/bits.py"
click node_onnx_io "https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/onnx_io.py"
click node_positions "https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/positions.py"
click node_selection "https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/selection.py"
click node_empirical "https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/empirical.py"
click node_build_proof "https://github.com/x-3306/onnxstego/blob/main/scripts/build_proof.py"
click node_demo_script "https://github.com/x-3306/onnxstego/blob/main/scripts/empirical_demo.py"
click node_proof_bundle "https://github.com/x-3306/onnxstego/blob/main/proof/README.md"
click node_test_bits "https://github.com/x-3306/onnxstego/blob/main/tests/test_bits.py"
click node_test_crypto "https://github.com/x-3306/onnxstego/blob/main/tests/test_crypto.py"
click node_test_onnx "https://github.com/x-3306/onnxstego/blob/main/tests/test_onnx_io.py"
click node_test_positions "https://github.com/x-3306/onnxstego/blob/main/tests/test_positions.py"
click node_test_selection "https://github.com/x-3306/onnxstego/blob/main/tests/test_selection.py"
click node_test_roundtrip "https://github.com/x-3306/onnxstego/blob/main/tests/test_core_roundtrip.py"
click node_test_empirical "https://github.com/x-3306/onnxstego/blob/main/tests/test_empirical.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_main,node_cli,node_core,node_crypto,node_bits,node_onnx_io,node_positions,node_selection,node_empirical toneBlue
class node_build_proof,node_demo_script toneAmber
class node_proof_bundle toneMint
class node_test_bits,node_test_crypto,node_test_onnx,node_test_positions,node_test_selection,node_test_roundtrip,node_test_empirical toneRose
```

## Components

### onnxstego package

Core library implementing authenticated message embedding and extraction in ONNX weights.

- [Module entry](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/__main__.py): entrypoint · Python module entry that launches the CLI. · `src/onnxstego/__main__.py`
- [CLI](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/cli.py): command interface · User-facing operations such as key generation, inspect, embed, and extract. · `src/onnxstego/cli.py`
- [Core flow](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/core.py): orchestrator · Coordinates encryption, bit packing, position selection, ONNX updates, and reverse extraction. · `src/onnxstego/core.py`
- [Crypto](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/crypto.py): aead layer · Authenticated encryption and decryption of payloads, with tamper and wrong-key rejection. · `src/onnxstego/crypto.py`
- [Bit codec](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/bits.py): bit packing · Converts bytes to bit streams and reconstructs bytes during extraction. · `src/onnxstego/bits.py`
- [ONNX I/O](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/onnx_io.py): tensor adapter · Reads and writes ONNX tensors while preserving valid float32 model serialization. · `src/onnxstego/onnx_io.py`
- [Positions](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/positions.py): keyed sampler · Derives deterministic embedding positions from the master key using a CSPRNG. · `src/onnxstego/positions.py`
- [Selection](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/selection.py): carrier filter · Chooses eligible weight locations, including uniform and natural reference-based modes. · `src/onnxstego/selection.py`
- [Empirical](https://github.com/x-3306/onnxstego/blob/main/src/onnxstego/empirical.py): analysis module · Measurement and smoke-test support for demonstration and validation workflows. · `src/onnxstego/empirical.py`

### Tooling

Command-line entrypoints and demo scripts.

- [Build proof](https://github.com/x-3306/onnxstego/blob/main/scripts/build_proof.py): demo script · Reconstructs the public proof assets and manifests. · `scripts/build_proof.py`
- [Demo script](https://github.com/x-3306/onnxstego/blob/main/scripts/empirical_demo.py): measurement script · Runs the empirical demonstration and basic measurements. · `scripts/empirical_demo.py`

### Proof bundle

Public demonstration assets and reproducibility manifest.

- [Proof assets](https://github.com/x-3306/onnxstego/blob/main/proof/README.md): artifact bundle · Reference, cover, and stego ONNX models plus a public manifest. · `proof/README.md`

### Tests

Subsystem-focused validation of codec, selection, serialization, and round-trip behavior.

- [Bits tests](https://github.com/x-3306/onnxstego/blob/main/tests/test_bits.py): test suite · Validates bit packing and unpacking behavior. · `tests/test_bits.py`
- [Crypto tests](https://github.com/x-3306/onnxstego/blob/main/tests/test_crypto.py): test suite · Checks authenticated encryption and rejection behavior. · `tests/test_crypto.py`
- [ONNX tests](https://github.com/x-3306/onnxstego/blob/main/tests/test_onnx_io.py): test suite · Validates ONNX tensor reading and writing behavior. · `tests/test_onnx_io.py`
- [Position tests](https://github.com/x-3306/onnxstego/blob/main/tests/test_positions.py): test suite · Checks deterministic keyed position derivation. · `tests/test_positions.py`
- [Selection tests](https://github.com/x-3306/onnxstego/blob/main/tests/test_selection.py): test suite · Validates uniform and natural carrier selection logic. · `tests/test_selection.py`
- [Round-trip](https://github.com/x-3306/onnxstego/blob/main/tests/test_core_roundtrip.py): end-to-end test · Exercises the full embed and extract pipeline. · `tests/test_core_roundtrip.py`
- [Empirical tests](https://github.com/x-3306/onnxstego/blob/main/tests/test_empirical.py): analysis test · Checks the empirical and statistical analysis path. · `tests/test_empirical.py`

## Connections

- Module entry → CLI: launches · Module execution enters the command-line interface.
- CLI → Core flow: invokes · CLI commands delegate to the orchestration layer.
- Core flow → Crypto: encrypts · Payload is authenticated before embedding and verified on extraction.
- Core flow → Bit codec: packs · Ciphertext is converted to and from bit streams.
- Core flow → ONNX I/O: reads/writes · Core flow mutates and serializes ONNX tensors.
- Core flow → Positions: derives · Keyed position derivation supplies embedding locations.
- Core flow → Selection: filters · Carrier eligibility and mode-specific selection shape where bits can land.
- Selection → ONNX I/O: depends on · Natural selection uses a reference ONNX model as part of carrier filtering.
- Positions → Selection: feeds · Deterministic positions are applied over the selected carrier set.
- Build proof → Proof assets: produces · Proof script reconstructs the public demonstration artifacts.
- Demo script → Empirical: runs · Demo script exercises the empirical analysis module.
- Empirical → Core flow: measures · Analysis code evaluates the embedding pipeline and its effects.
- Bits tests → Bit codec: covers · Bit codec tests validate stream conversion logic.
- Crypto tests → Crypto: covers · Crypto tests validate AEAD correctness and rejection.
- ONNX tests → ONNX I/O: covers · Serialization tests validate ONNX tensor handling.
- Position tests → Positions: covers · Position tests validate keyed sampling behavior.
- Selection tests → Selection: covers · Selection tests validate carrier filtering modes.
- Round-trip → Core flow: covers · End-to-end tests exercise the full embed/extract pipeline.
- Empirical tests → Empirical: covers · Empirical tests validate measurement and statistical checks.
- Proof assets → ONNX I/O: demonstrates · Public artifacts demonstrate the ONNX carrier format in practice.