Pure Python 3 implementation of Rust `rand`'s StdRng random number generator.

rust_stdrng.py

"""
Pure Python 3 implementation of the random number generator
used in Rust's `rand` crate, particularly StdRng.

This code implements:

- The PCG32 pseudorandom number generator used by `SeedableRng::seed_from_u64`
  in order to seed the underlying RNG in turn.

  https://docs.rs/rand_core/0.6.4/rand_core/trait.SeedableRng.html#method.seed_from_u64

- A configurable-round ChaCha keystream generator.

- Some convenience functions that allow to use the entire combination to implement
  Rust's `rand` crate's StdRng: https://docs.rs/rand/latest/rand/rngs/struct.StdRng.html.
  It may also be used to implement any of `rand_chacha`'s PRNGs by configuring the
  correct number of rounds (8, 12, 20).

Barely tested, was put together as part of the following CTF writeup:

https://nnub.es/blog/en/ctf/corctf/2024/cormine

Adapts the PCG32 Python generator implementation from Rosetta Code:
https://rosettacode.org/wiki/Pseudo-random_numbers/PCG32#Python:_As_generator

ChaCha generator adapted from Péter Szabó's pure Python implementation for ChaCha20:
https://github.com/pts/chacha20/blob/8a4b31641daf42efd173117c8fbe847f7c0e6b32/chacha20_python3.py

- Based on https://gist.github.com/cathalgarvey/0ce7dbae2aa9e3984adc
- Based on Numpy implementation: https://gist.github.com/chiiph/6855750
- Based on http://cr.yp.to/chacha.html

More info about ChaCha20: https://en.wikipedia.org/wiki/Salsa20

IMPORTANT NOTE: No license is explicitly provided, as no license was attached to
the original ChaCha20 code by Péter Szabó. Use this code at your own risk.
"""

from itertools import islice
import struct


def pcg32(seed_state: int = None, inc: int = None, as_int=True):
    mask64 = (1 << 64) - 1
    mask32 = (1 << 32) - 1
    CONST = 6364136223846793005

    def next_int():
        "return random 32 bit unsigned int"
        nonlocal state, inc

        state = ((state * CONST) + inc) & mask64
        xorshifted, rot = (
            (((state >> 18) ^ state) >> 27) & mask32,
            (state >> 59) & mask32,
        )
        answer = ((xorshifted >> rot) | (xorshifted << ((-rot) & 31))) & mask32
        return answer

    state = seed_state

    while True:
        yield next_int() if as_int else next_int() / (1 << 32)


def gen_chacha_seed(base_seed: int):
    chacha_seed_length = 8  # 32 bytes / 4 (sizeof(uint32_t)) = 8 integers
    chacha_seed = bytearray()
    for k in islice(pcg32(base_seed, 11634580027462260723, True), chacha_seed_length):
        k_bytes = struct.pack("I", k)
        chacha_seed.extend(k_bytes)
    return bytes(chacha_seed)


def chacha_xor_stream(key, iv=(b"\x00" * 8), position=0, rounds=20):
    """Generate the xor stream with the ChaCha20 cipher."""
    if not isinstance(position, int):
        raise TypeError
    if position & ~0xFFFFFFFF:
        raise ValueError("Position is not uint32.")
    if not isinstance(key, bytes):
        raise TypeError
    if not isinstance(iv, bytes):
        raise TypeError
    if len(key) != 32:
        raise ValueError
    if len(iv) != 8:
        raise ValueError
    if not isinstance(rounds, int):  # NEW
        raise TypeError
    if rounds % 2 != 0:  # NEW
        raise ValueError(f"Number of rounds needs to be even, got: {rounds}")

    # NEW: For convenience we'll assume we only need double rounds.
    double_rounds = int(rounds / 2)

    def rotate(v, c):
        return ((v << c) & 0xFFFFFFFF) | v >> (32 - c)

    def quarter_round(x, a, b, c, d):
        x[a] = (x[a] + x[b]) & 0xFFFFFFFF
        x[d] = rotate(x[d] ^ x[a], 16)
        x[c] = (x[c] + x[d]) & 0xFFFFFFFF
        x[b] = rotate(x[b] ^ x[c], 12)
        x[a] = (x[a] + x[b]) & 0xFFFFFFFF
        x[d] = rotate(x[d] ^ x[a], 8)
        x[c] = (x[c] + x[d]) & 0xFFFFFFFF
        x[b] = rotate(x[b] ^ x[c], 7)

    ctx = [0] * 16
    ctx[:4] = (1634760805, 857760878, 2036477234, 1797285236)
    ctx[4:12] = struct.unpack("<8L", key)
    ctx[12] = ctx[13] = position
    ctx[14:16] = struct.unpack("<LL", iv)
    while 1:
        x = list(ctx)
        for _ in range(double_rounds):  # CHANGED: dynamic amount of double rounds
            quarter_round(x, 0, 4, 8, 12)
            quarter_round(x, 1, 5, 9, 13)
            quarter_round(x, 2, 6, 10, 14)
            quarter_round(x, 3, 7, 11, 15)
            quarter_round(x, 0, 5, 10, 15)
            quarter_round(x, 1, 6, 11, 12)
            quarter_round(x, 2, 7, 8, 13)
            quarter_round(x, 3, 4, 9, 14)
        for c in struct.pack(
            "<16L", *((x[i] + ctx[i]) & 0xFFFFFFFF for i in range(16))
        ):
            yield c
        ctx[12] = (ctx[12] + 1) & 0xFFFFFFFF
        if ctx[12] == 0:
            ctx[13] = (ctx[13] + 1) & 0xFFFFFFFF


def rust_stdrng_chacha(seed, iv=(b"\x00" * 8), position=0, rounds=12):
    yield from chacha_xor_stream(
        gen_chacha_seed(seed), iv=iv, position=position, rounds=rounds
    )


def __usage_example():
    seed = 123456789  # u64 seed
    keystream_iter = rust_stdrng_chacha(seed, rounds=12)
    # Now use keystream_iter as an infinite iterator for keystream bytes.
    # Example usage as a stream cipher:
    plaintext = b"Hello World!"
    ciphertext = bytearray()
    print(plaintext)
    for k, p in zip(keystream_iter, plaintext):
        ciphertext.append(k ^ p)
    print(bytes(ciphertext))
    keystream_iter = rust_stdrng_chacha(seed, rounds=12)  # reset for decryption
    plaintext = bytearray()
    for k, p in zip(keystream_iter, ciphertext):
        plaintext.append(k ^ p)
    print(bytes(plaintext))


if __name__ == "__main__":
    __usage_example()