#!/usr/bin/env python3
"""
Bug #15 — PoC: RNG::nextInt Modular Bias and Division-by-Zero

Vulnerability
=============
RNG::nextInt(uint32_t _min, uint32_t _max) in Utils.cpp:

    uint32_t num;
    random((uint8_t *) &num, sizeof(num));
    return (num % (_max - _min)) + _min;

Two defects:

1. **Division-by-zero (undefined behavior):**
   If _max == _min, the expression becomes `num % 0`. In C++ this is
   undefined behavior. On ARM Cortex-M (nRF52, ESP32) there is no
   hardware divide-by-zero trap for unsigned integer modulo —
   behavior depends on compiler/runtime: GCC on ARM typically returns 0
   but behavior varies. Some toolchains do trap. The real risk is a
   future caller passing equal bounds (defensive code should not rely
   on "current callers are safe").

2. **Modular bias:**
   `num % range` produces a biased distribution when `range` does not
   evenly divide 2^32. For small ranges (< 100) the bias is negligible.
   For the BLE PIN range (100000..999999, range = 899999), the bias
   reduces the effective entropy:

     2^32 mod 899999 = 294,969
     → Values 0..294,968 are 1 more likely than 294,969..899,998
     → Bias: 32.8% of PIN space gets +1/4773 extra weight
     → Reduces effective entropy from ~19.78 bits to ~19.77 bits

   In practice this is a negligible security impact for BLE PINs.

Attack surface
--------------
- **Div-by-zero:** Not remotely triggerable in current firmware (all
  callers use hardcoded bounds where _min < _max). Risk is latent —
  a future caller could pass computed bounds where equality is possible.

- **Modular bias:** Affects BLE PIN generation (nextInt(100000, 999999)).
  The bias is ~0.005% — negligible for a 6-digit PIN that's also shown
  on the device display. The real BLE PIN concern is that the range
  [100000, 899999) excludes PINs 900000–999999 (likely a bug: should be
  nextInt(100000, 1000000) for full 6-digit range).

Affected code
-------------
  src/Utils.h   — RNG class declaration
  src/Utils.cpp — nextInt implementation
  examples/companion_radio/MyMesh.cpp — BLE PIN generation

This PoC:
  Part A — Demonstrates modular bias visually
  Part B — Shows div-by-zero UB scenario
  Part C — Analyzes BLE PIN range bias + excluded values
"""

import os
import struct
import collections

# ──────────────────────────────────────────────────────────────────────
# Simulate RNG::nextInt in Python
# ──────────────────────────────────────────────────────────────────────
def nextInt_biased(_min, _max):
    """Exact replica of the firmware's nextInt — modular bias included."""
    num = struct.unpack("<I", os.urandom(4))[0]
    rang = _max - _min
    if rang == 0:
        raise ZeroDivisionError("_max == _min → num % 0 → undefined behavior!")
    return (num % rang) + _min

def nextInt_unbiased(_min, _max):
    """Rejection-sampling method — no bias."""
    rang = _max - _min
    if rang == 0:
        return _min
    threshold = (0x100000000 - rang) % rang  # = 2^32 mod range
    while True:
        num = struct.unpack("<I", os.urandom(4))[0]
        if num >= threshold:
            return (num % rang) + _min


# ══════════════════════════════════════════════════════════════════════
# Part A — Modular bias visualization (small range)
# ══════════════════════════════════════════════════════════════════════
print("=" * 70)
print("Part A: Modular bias demonstration (range = 3)")
print("=" * 70)

# For range = 3: 2^32 = 1,431,655,765 * 3 + 1
# → value 0 gets 1,431,655,766 hits, values 1-2 get 1,431,655,765
# → value 0 is 0.0000000233% more likely
# Use a much smaller simulated space to exaggerate the effect
SIMULATED_BITS = 8  # simulate with 8-bit random for visible bias
SIMULATED_MAX = 2**SIMULATED_BITS  # 256
RANGE = 3
TRIALS = 1_000_000

counts_biased = collections.Counter()
counts_uniform = collections.Counter()

for _ in range(TRIALS):
    num_biased = struct.unpack("B", os.urandom(1))[0]  # 0..255
    counts_biased[num_biased % RANGE] += 1

    # Rejection sampling (unbiased)
    while True:
        num_unbiased = struct.unpack("B", os.urandom(1))[0]
        if num_unbiased >= (SIMULATED_MAX % RANGE):  # threshold = 256 % 3 = 1
            counts_uniform[num_unbiased % RANGE] += 1
            break

print(f"  Using 8-bit random (0..255) with range = {RANGE}")
print(f"  256 mod 3 = {256 % 3} → value 0 gets {256 // 3 + 1} slots, values 1-2 get {256 // 3} slots")
print(f"  Expected bias: value 0 is {(256//3+1)/(256//3) - 1:.2%} more likely")
print(f"\n  Biased (firmware method):    {dict(sorted(counts_biased.items()))}")
print(f"  Unbiased (rejection method): {dict(sorted(counts_uniform.items()))}")
bias_pct = (counts_biased[0] / TRIALS - 1/3) / (1/3) * 100
print(f"  Observed bias for value 0: {bias_pct:+.2f}%")


# ══════════════════════════════════════════════════════════════════════
# Part B — Division-by-zero undefined behavior
# ══════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("Part B: Division-by-zero when _min == _max")
print("=" * 70)

print("  Calling nextInt(5, 5) — firmware equivalent: num % 0")
try:
    result = nextInt_biased(5, 5)
    print(f"  Result: {result} (Python raises, C++ is UB → platform-dependent)")
except ZeroDivisionError as e:
    print(f"  ✗ ZeroDivisionError: {e}")
    print(f"    On ARM Cortex-M4 (nRF52): behavior is undefined")
    print(f"    GCC on ARM typically returns 0 for unsigned div-by-zero")
    print(f"    Some toolchains trap → HardFault → device reboot")

print("\n  Calling nextInt(0, 0):")
try:
    result = nextInt_biased(0, 0)
    print(f"  Result: {result}")
except ZeroDivisionError as e:
    print(f"  ✗ ZeroDivisionError: {e}")

print("\n  All current firmware callers (safe — hardcoded bounds):")
callers = [
    ("Mesh.cpp getRetransmitDelay",        0, 5,   "retransmit jitter"),
    ("Mesh.cpp getCADFailRetryDelay",       1, 4,   "CAD retry jitter"),
    ("SensorMesh.cpp getRetransmitDelay",   0, 6,   "sensor retransmit"),
    ("MyMesh.cpp getRetransmitDelay",       0, "5*t+1", "repeater retransmit"),
    ("MyMesh.cpp BLE PIN",            100000, 999999, "BLE pairing PIN"),
]
for name, mn, mx, purpose in callers:
    if isinstance(mx, int):
        safe = "✓ safe" if mn != mx else "✗ DIVZERO"
        print(f"    {name}: nextInt({mn}, {mx}) → {safe} — {purpose}")
    else:
        print(f"    {name}: nextInt({mn}, {mx}) → ✓ safe (t≥0 → max≥1) — {purpose}")

print("\n  Risk: A FUTURE caller could pass computed bounds where min==max")
print("  Example: nextInt(0, num_peers) when num_peers == 0 → crash")


# ══════════════════════════════════════════════════════════════════════
# Part C — BLE PIN range analysis
# ══════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("Part C: BLE PIN range bias analysis")
print("=" * 70)

PIN_MIN = 100000
PIN_MAX = 999999  # exclusive in firmware: nextInt(100000, 999999)
PIN_RANGE = PIN_MAX - PIN_MIN  # = 899999

remainder = (2**32) % PIN_RANGE
full_cycles = (2**32) // PIN_RANGE

print(f"  Firmware: nextInt({PIN_MIN}, {PIN_MAX})")
print(f"  Effective range: [{PIN_MIN}, {PIN_MAX}) → {PIN_RANGE} values")
print(f"  NOTE: PINs 999999 is excluded! Range should likely be [100000, 1000000)")
print(f"\n  Modular bias analysis:")
print(f"    2^32 = {full_cycles} × {PIN_RANGE} + {remainder}")
print(f"    Values 0..{remainder-1}: appear {full_cycles + 1} times → PINs {PIN_MIN}..{PIN_MIN + remainder - 1}")
print(f"    Values {remainder}..{PIN_RANGE-1}: appear {full_cycles} times → PINs {PIN_MIN + remainder}..{PIN_MAX - 1}")
print(f"    Overrepresented PINs: {remainder} / {PIN_RANGE} = {remainder/PIN_RANGE*100:.2f}%")
print(f"    Bias magnitude: {1/(full_cycles+1)*100 - 1/(full_cycles+1)*100:.6f}% per overrepresented PIN")

import math
# Effective entropy
ideal_entropy = math.log2(PIN_RANGE)
# With bias: p_high = (full_cycles+1)/2^32, p_low = full_cycles/2^32
p_high = (full_cycles + 1) / 2**32
p_low = full_cycles / 2**32
entropy = -(remainder * p_high * math.log2(p_high) + (PIN_RANGE - remainder) * p_low * math.log2(p_low))
print(f"\n  Ideal entropy:    {ideal_entropy:.4f} bits")
print(f"  Actual entropy:   {entropy:.4f} bits")
print(f"  Entropy loss:     {ideal_entropy - entropy:.6f} bits (negligible)")


# ══════════════════════════════════════════════════════════════════════
# Summary
# ══════════════════════════════════════════════════════════════════════
print("\n" + "=" * 70)
print("SUMMARY")
print("=" * 70)
print("""
Two defects in RNG::nextInt():

1. Division-by-zero (UB) if _max == _min:
   - Not triggerable in current firmware (all bounds are hardcoded safe)
   - Latent risk for future callers
   - Fix: guard _max <= _min → return _min

2. Modular bias from `num % range`:
   - For small ranges (retransmit jitter): bias < 0.0001% → irrelevant
   - For BLE PIN (range=899999): ~0.005% bias → negligible
   - BLE PIN range bug: [100000, 999999) excludes 999999
   - Fix: rejection sampling to eliminate bias

Severity: LOW
  - No remotely triggerable crash path in current firmware
  - Bias is mathematically present but practically insignificant
  - BLE PIN excluded-range bug is minor (reduces space by 1 value)
  - Main value is defensive hardening against future misuse
""")