Python random numbers from KMS.GenerateRandom

kms_random.md

Python random numbers from KMS.GenerateRandom

Spurred by this twitter conversation. random.SystemRandom uses os.urandom as a source of bytes, but doesn't provide a way to use a different source of bytes. So stream_random.py is exactly that. Then kms_random.py has raw and buffered bytestreams pulling from KMS.GenerateRandom.

The main interface is kms_random.get_kms_random(boto3_session, buffer_size=None). The default buffer size is 16, chosen arbitrarily.

I do not vouch for the randomness properties of the results. Presumably using the same logic as SystemRandom should work, but I am not an expert in cryptography or randomness.

kms_random.py

# MIT No Attribution
#
# Copyright 2023 Ben Kehoe
#
# Permission is hereby granted, free of charge, to any person obtaining a copy of this
# software and associated documentation files (the "Software"), to deal in the Software
# without restriction, including without limitation the rights to use, copy, modify,
# merge, publish, distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A
# PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
# OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

import io
import random

import boto3

from stream_random import BytestreamRandom

__all__ = ["get_kms_random"]

def get_kms_random(session: boto3.Session, buffer_size=None) -> random.Random:
    """Get a random number generator using KMS.GenerateRandom as a source"""
    kms_bytestream = KMSByteStream(session, buffer_size=buffer_size)
    return BytestreamRandom(kms_bytestream)

class KMSRawByteStream(io.RawIOBase):
    """Raw bytestream from KMS.GenerateRandom"""

    def __init__(self, session: boto3.Session) -> None:
        super().__init__()
        self.session = session
        self.kms_client = session.client("kms")
    
    def _get_bytes(self, num_bytes: int) -> bytes:
        response = self.kms_client.generate_random(NumberOfBytes=num_bytes)
        return response["Plaintext"]
    
    def _check_closed(self):
        if self.closed:
            raise ValueError("Closed")

    def readable(self) -> bool:
        return True
    
    def readall(self) -> bytes:
        raise io.UnsupportedOperation("readall")
    
    def readinto(self, __buffer) -> int | None:
        self._check_closed()
        num_bytes = len(__buffer)
        bytes = self._get_bytes(num_bytes)
        __buffer[:len(bytes)] = bytes
        return len(bytes)

    # seek

    def seekable(self) -> bool:
        return False
    
    # tell, truncate

    def writable(self) -> bool:
        return False
    
    def write(self, __b) -> int | None:
        raise io.UnsupportedOperation("write")

class KMSByteStream(io.BufferedReader):
    """Buffered bytestream from KMS.GenerateRandom"""
    def __init__(self, session: boto3.Session, buffer_size=None) -> None:
        raw = KMSRawByteStream(session)
        if buffer_size is None:
            buffer_size = 16 # UUID-sized
        super().__init__(raw, buffer_size=buffer_size)

stream_random.py

# Modified from stdlib random.SystemRandom
# PSF License
# https://docs.python.org/3/license.html#psf-license

import random
from io import BufferedIOBase

BPF = 53        # Number of bits in a float
RECIP_BPF = 2 ** -BPF

class BytestreamRandom(random.Random):
    """random number generator using a stream of bytes as a source"""

    def __init__(self, binary_stream: BufferedIOBase):
        self.stream = binary_stream
    
    def _get_bytes(self, num_bytes):
        return self.stream.read(num_bytes)

    def random(self):
        """Get the next random number in the range 0.0 <= X < 1.0."""
        return (int.from_bytes(self._get_bytes(7), byteorder='big') >> 3) * RECIP_BPF

    def getrandbits(self, k):
        """getrandbits(k) -> x.  Generates an int with k random bits."""
        if k < 0:
            raise ValueError('number of bits must be non-negative')
        numbytes = (k + 7) // 8                       # bits / 8 and rounded up
        x = int.from_bytes(self._get_bytes(numbytes), byteorder='big')
        return x >> (numbytes * 8 - k)                # trim excess bits

    def randbytes(self, n):
        """Generate n random bytes."""
        return self._get_bytes(n)

    def seed(self, *args, **kwds):
        "Stub method.  Not used for this random number generator."
        return None

    def _notimplemented(self, *args, **kwds):
        "Method should not be called for a this random number generator."
        raise NotImplementedError('Source does not have state.')
    getstate = setstate = _notimplemented