From 7888d7fc73cefea338dea4c62d5354c5422256db Mon Sep 17 00:00:00 2001
From: highlander <highlander@keepkey.com>
Date: Mon, 3 Aug 2026 14:06:44 -0300
Subject: [PATCH] test(rng): prove entropy audit budget policy

---
 scripts/generate-test-report.py |  8 +++-
 tests/test_msg_getentropy.py    | 76 ++++++++++++++++++++++++---------
 tests/test_protection_levels.py |  1 +
 3 files changed, 64 insertions(+), 21 deletions(-)

diff --git a/scripts/generate-test-report.py b/scripts/generate-test-report.py
index 70056b87..a80289bc 100644
--- a/scripts/generate-test-report.py
+++ b/scripts/generate-test-report.py
@@ -644,8 +644,12 @@ def _arg_shown(a):
           'or information leaks. Verifies input sanitization.',
           []),
          ('C27', 'test_msg_getentropy', 'test_entropy',
-          'Hardware RNG entropy',
-          'Reads random bytes from the hardware RNG. Used to verify the entropy source is functional.',
+          'Hardware RNG audit budget and lock policy',
+          'Proves a fresh initialized, PIN-protected, locked device still requires confirmation; '
+          'then proves an uninitialized device returns exactly 8 x 8192 bytes (64 KiB) without a '
+          'press, with exact lengths, unique blocks, and conservative catastrophic-failure health '
+          'checks. The next request must restore confirmation. These checks detect a stuck or '
+          'grossly biased source; they are not a statistical certification of the hardware RNG.',
           []),
          ('C28', 'test_msg_cipherkeyvalue', 'test_encrypt',
           'Symmetric key encryption',
diff --git a/tests/test_msg_getentropy.py b/tests/test_msg_getentropy.py
index 96ea7abe..f12d4f90 100644
--- a/tests/test_msg_getentropy.py
+++ b/tests/test_msg_getentropy.py
@@ -20,35 +20,73 @@
 
 from __future__ import print_function
 
+import os
 import unittest
 import common
-import math
+from collections import Counter
 
 import keepkeylib.messages_pb2 as proto
 import keepkeylib.types_pb2 as proto_types
 
-def entropy(data):
-    counts = {}
-    for c in data:
-        if c in counts:
-            counts[c] += 1
-        else:
-            counts[c] = 1
-    e = 0
-    for _, v in counts.items():
-        p = 1.0 * v / len(data)
-        e -= p * math.log(p, 256)
-    return e
-
 class TestMsgGetentropy(common.KeepKeyTest):
 
+    @unittest.skipUnless(
+        os.getenv('KK_EXPECT_ENTROPY_BUDGET') == '1',
+        'requires the RC23 entropy audit budget policy')
     def test_entropy(self):
-        for l in [0, 1, 2, 3, 4, 5, 8, 9, 16, 17, 32, 33, 64, 65, 128, 129, 256, 257, 512, 513, 1024]:
+        chunk_size = 8192
+        chunk_count = 8
+
+        # A fresh budget must not make raw RNG output silently available from
+        # an initialized, PIN-protected, locked device.  Confirm one request in
+        # that state before spending any of the press-free budget.
+        self.setup_mnemonic_pin_passphrase()
+        self.client.clear_session()
+        with self.client:
+            self.client.set_expected_responses([
+                proto.ButtonRequest(code=proto_types.ButtonRequest_GetEntropy),
+                proto.Entropy(),
+            ])
+            locked_sample = self.client.get_entropy(chunk_size)
+        self.assertEqual(len(locked_sample), chunk_size)
+
+        # Wiping returns the device to the uninitialized audit state.  The
+        # confirmed locked request above does not consume the fresh budget.
+        self.client.wipe_device()
+
+        samples = []
+        for _ in range(chunk_count):
             with self.client:
-                self.client.set_expected_responses([proto.ButtonRequest(code=proto_types.ButtonRequest_GetEntropy), proto.Entropy()])
-                ent = self.client.get_entropy(l)
-                self.assertTrue(len(ent) >= l)
-                print('entropy = ', entropy(ent))
+                self.client.set_expected_responses([proto.Entropy()])
+                sample = self.client.get_entropy(chunk_size)
+            self.assertEqual(len(sample), chunk_size)
+            samples.append(sample)
+
+        self.assertEqual(sum(len(sample) for sample in samples), 64 * 1024)
+        self.assertEqual(len(set(samples)), chunk_count)
+
+        # Deliberately broad catastrophic-failure checks, not a statistical
+        # certification of the hardware RNG.  They catch a stuck/constant or
+        # grossly biased source without imposing a fragile quality threshold.
+        combined = b''.join(samples)
+        counts = Counter(combined)
+        self.assertGreaterEqual(len(counts), 200)
+        self.assertLess(max(counts.values()), len(combined) // 20)
+        one_bits = sum(bin(value).count('1') for value in combined)
+        one_ratio = float(one_bits) / (8 * len(combined))
+        self.assertGreater(one_ratio, 0.40)
+        self.assertLess(one_ratio, 0.60)
+
+        # Exactly 64 KiB was press-free.  The next request must restore the
+        # original confirmation flow and still return the requested length
+        # after the debug-link approval.
+        with self.client:
+            self.client.set_expected_responses([
+                proto.ButtonRequest(code=proto_types.ButtonRequest_GetEntropy),
+                proto.Entropy(),
+            ])
+            after_budget = self.client.get_entropy(chunk_size)
+        self.assertEqual(len(after_budget), chunk_size)
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/test_protection_levels.py b/tests/test_protection_levels.py
index 2efd5676..a9fc6638 100644
--- a/tests/test_protection_levels.py
+++ b/tests/test_protection_levels.py
@@ -68,6 +68,7 @@ def test_ping(self):
     def test_get_entropy(self):
         with self.client:
             self.setup_mnemonic_pin_passphrase()
+            self.client.clear_session()
             self.client.set_expected_responses([proto.ButtonRequest(),
                                       proto.Entropy()])
             self.client.get_entropy(10)