Skip to content

Instantly share code, notes, and snippets.

@benkehoe
Created April 28, 2023 14:45
Show Gist options
  • Save benkehoe/cc4768020c3f0454c59d2dc1d33e4fa2 to your computer and use it in GitHub Desktop.
Save benkehoe/cc4768020c3f0454c59d2dc1d33e4fa2 to your computer and use it in GitHub Desktop.
Python random numbers from KMS.GenerateRandom

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.

# 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)
# 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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment