test_agent_seed_zeroing.py
python
sha256:88ac91129873e6a496e9189515aa690eb893ae25d69c8f72af141a2be5068eb3
docs: docstring sprint contract→find-symbol — idiomatic run…
Sonnet 4.6
patch
140 days ago
| 1 | """Tests for memory zeroing in muse/cli/commands/agent.py key derivation. |
| 2 | |
| 3 | Two functions in agent.py create DerivedKey objects without zeroing them: |
| 4 | |
| 5 | 1. _derive_agent_seed: calls mnemonic_to_seed() → bytearray, passes it to |
| 6 | derive_agent_sub_seed(), but never zeros the intermediate seed bytearray. |
| 7 | |
| 8 | 2. _sub_seed_to_public: calls derive_identity_key(sub_seed) → DerivedKey, |
| 9 | creates the Ed25519 key, but never calls dk.zero() on the DerivedKey. |
| 10 | |
| 11 | Fix: |
| 12 | - _derive_agent_seed: zero the seed bytearray after derive_agent_sub_seed(). |
| 13 | - _sub_seed_to_public: call dk.zero() after creating the Ed25519 key object. |
| 14 | |
| 15 | Coverage |
| 16 | -------- |
| 17 | I _derive_agent_seed zeroes intermediate seed |
| 18 | I1 seed bytearray is zeroed after derive_agent_sub_seed call |
| 19 | |
| 20 | II _sub_seed_to_public zeroes DerivedKey |
| 21 | II1 DerivedKey.private_bytes is zeroed after _sub_seed_to_public returns |
| 22 | II2 DerivedKey.chain_code is zeroed after _sub_seed_to_public returns |
| 23 | """ |
| 24 | |
| 25 | from __future__ import annotations |
| 26 | |
| 27 | from unittest.mock import patch |
| 28 | |
| 29 | import pytest |
| 30 | |
| 31 | from muse.core import hdkeys as _hdkeys |
| 32 | from muse.core.bip39 import mnemonic_to_seed |
| 33 | from muse.core.hdkeys import derive_agent_sub_seed, DOMAIN_IDENTITY |
| 34 | from muse.core.slip010 import DerivedKey |
| 35 | import muse.core.bip39 as _bip39_mod |
| 36 | |
| 37 | _MNEMONIC = ( |
| 38 | "abandon abandon abandon abandon abandon abandon abandon abandon " |
| 39 | "abandon abandon abandon about" |
| 40 | ) |
| 41 | _SEED = mnemonic_to_seed(_MNEMONIC) |
| 42 | _SUB_SEED = derive_agent_sub_seed(_SEED, domain=DOMAIN_IDENTITY, agent_id=0) |
| 43 | |
| 44 | |
| 45 | # --------------------------------------------------------------------------- |
| 46 | # I _derive_agent_seed zeroes intermediate seed |
| 47 | # --------------------------------------------------------------------------- |
| 48 | |
| 49 | class TestDeriveAgentSeedZeroesIntermediateSeed: |
| 50 | def test_I1_seed_zeroed_after_derive_agent_sub_seed(self) -> None: |
| 51 | """I1: the bytearray seed from mnemonic_to_seed is zeroed after derive_agent_sub_seed.""" |
| 52 | from muse.cli.commands.agent import _derive_agent_seed |
| 53 | |
| 54 | seeds_captured: list[bytearray] = [] |
| 55 | original_to_seed = _bip39_mod.mnemonic_to_seed |
| 56 | |
| 57 | def tracking_to_seed(words, passphrase=""): |
| 58 | result = original_to_seed(words, passphrase) |
| 59 | seeds_captured.append(result) |
| 60 | return result |
| 61 | |
| 62 | with patch.object(_bip39_mod, "mnemonic_to_seed", side_effect=tracking_to_seed): |
| 63 | _derive_agent_seed(_MNEMONIC, account=0) |
| 64 | |
| 65 | assert seeds_captured, "mnemonic_to_seed was not called" |
| 66 | seed = seeds_captured[0] |
| 67 | assert seed == bytearray(64), ( |
| 68 | "intermediate seed bytearray must be zeroed after _derive_agent_seed returns" |
| 69 | ) |
| 70 | |
| 71 | |
| 72 | # --------------------------------------------------------------------------- |
| 73 | # II _sub_seed_to_public zeroes DerivedKey |
| 74 | # --------------------------------------------------------------------------- |
| 75 | |
| 76 | class TestSubSeedToPublicZeroesDerivedKey: |
| 77 | def test_II1_private_bytes_zeroed(self) -> None: |
| 78 | """II1: DerivedKey.private_bytes is zeroed after _sub_seed_to_public returns.""" |
| 79 | from muse.cli.commands.agent import _sub_seed_to_public |
| 80 | |
| 81 | captured: list[DerivedKey] = [] |
| 82 | original_derive = _hdkeys.derive_identity_key |
| 83 | |
| 84 | def capturing_derive(*args, **kwargs): |
| 85 | dk = original_derive(*args, **kwargs) |
| 86 | captured.append(dk) |
| 87 | return dk |
| 88 | |
| 89 | with patch.object(_hdkeys, "derive_identity_key", side_effect=capturing_derive): |
| 90 | _sub_seed_to_public(bytes(_SUB_SEED)) |
| 91 | |
| 92 | assert captured, "derive_identity_key was not called" |
| 93 | dk = captured[0] |
| 94 | assert dk.private_bytes == bytearray(32), ( |
| 95 | "DerivedKey.private_bytes must be zeroed after _sub_seed_to_public" |
| 96 | ) |
| 97 | |
| 98 | def test_II2_chain_code_zeroed(self) -> None: |
| 99 | """II2: DerivedKey.chain_code is zeroed after _sub_seed_to_public returns.""" |
| 100 | from muse.cli.commands.agent import _sub_seed_to_public |
| 101 | |
| 102 | captured: list[DerivedKey] = [] |
| 103 | original_derive = _hdkeys.derive_identity_key |
| 104 | |
| 105 | def capturing_derive(*args, **kwargs): |
| 106 | dk = original_derive(*args, **kwargs) |
| 107 | captured.append(dk) |
| 108 | return dk |
| 109 | |
| 110 | with patch.object(_hdkeys, "derive_identity_key", side_effect=capturing_derive): |
| 111 | _sub_seed_to_public(bytes(_SUB_SEED)) |
| 112 | |
| 113 | assert captured, "derive_identity_key was not called" |
| 114 | dk = captured[0] |
| 115 | assert dk.chain_code == bytearray(32), ( |
| 116 | "DerivedKey.chain_code must be zeroed after _sub_seed_to_public" |
| 117 | ) |
File History
1 commit
sha256:88ac91129873e6a496e9189515aa690eb893ae25d69c8f72af141a2be5068eb3
docs: docstring sprint contract→find-symbol — idiomatic run…
Sonnet 4.6
patch
140 days ago