diff --git a/refinery/lib/argformats.py b/refinery/lib/argformats.py index 9f27da304..817ccb953 100644 --- a/refinery/lib/argformats.py +++ b/refinery/lib/argformats.py @@ -1135,6 +1135,55 @@ def cycle(self, it: Iterable[int]) -> Iterable[int]: return it return itertools.cycle(it) + @handler.register('rng', final=True) + def rng( + self, + size: str, + ) -> bytes: + """ + The `rng:count` handler generates `count` secure random bytes using the Python standard + library module function `secrets.token_bytes`. + """ + import secrets + size = PythonExpression.Lazy(size) + try: + _size = size() + except ParserVariableMissing: + def finalize(data): + meta = dict(metavars(data)) + return secrets.token_bytes(size(meta)) + return finalize + else: + return secrets.token_bytes(_size) + + @handler.register('prng', final=True) + def prng( + self, + size: str, + seed: Optional[str] = None + ) -> bytes: + """ + The `prng[seed]:count` handler generates `count` random bytes using Python's built-in random + number generator. The `seed` argument can be omitted, in which case the PRNG will be seeded + with the current timestamp in nanoseconds. + """ + import random + import time + seed = time.time_ns if seed is None else PythonExpression.Lazy(seed) + size = PythonExpression.Lazy(size) + try: + _size = size() + _seed = seed() + except ParserVariableMissing: + def finalize(data): + meta = dict(metavars(data)) + random.seed(seed(meta)) + return random.randbytes(size(meta)) + return finalize + else: + random.seed(_seed) + return random.randbytes(_size) + @handler.register('accu', final=True) def accu( self, diff --git a/test/lib/test_argformats.py b/test/lib/test_argformats.py index 87c3a59dc..309146c0d 100644 --- a/test/lib/test_argformats.py +++ b/test/lib/test_argformats.py @@ -3,6 +3,7 @@ from argparse import ArgumentTypeError from refinery.lib import argformats +from refinery.lib.tools import entropy from .. import TestBase @@ -164,3 +165,14 @@ def test_read_handler(self): os.unlink(name) except Exception: pass + + def test_prng(self): + import random + random.seed(0x1337) + goal = random.randbytes(2000) + test = argformats.multibin('prng[0x1337]:2000') + self.assertEqual(goal, test) + + def test_rng(self): + test = argformats.multibin('rng:200000') + self.assertGreaterEqual(entropy(test), 0.99)