adds random number generator handlers

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,

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)