1  from random import Random
  2  from typing import Any
  3  from enum import Enum
  4  
  5  
  6  UINT_SIZES = [8, 16, 32, 64, 128, 256]
  7  
  8  basic_types = ["uint%d" % v for v in UINT_SIZES] + ['bool', 'byte']
  9  
 10  random_mode_names = ["random", "zero", "max", "nil", "one", "lengthy"]
 11  
 12  
 13  class RandomizationMode(Enum):
 14      # random content / length
 15      mode_random = 0
 16      # Zero-value
 17      mode_zero = 1
 18      # Maximum value, limited to count 1 however
 19      mode_max = 2
 20      # Return 0 values, i.e. empty
 21      mode_nil_count = 3
 22      # Return 1 value, random content
 23      mode_one_count = 4
 24      # Return max amount of values, random content
 25      mode_max_count = 5
 26  
 27      def to_name(self):
 28          return random_mode_names[self.value]
 29  
 30      def is_changing(self):
 31          return self.value in [0, 4, 5]
 32  
 33  
 34  def get_random_ssz_object(rng: Random, typ: Any, max_bytes_length: int, max_list_length: int, mode: RandomizationMode, chaos: bool) -> Any:
 35      """
 36      Create an object for a given type, filled with random data.
 37      :param rng: The random number generator to use.
 38      :param typ: The type to instantiate
 39      :param max_bytes_length: the max. length for a random bytes array
 40      :param max_list_length: the max. length for a random list
 41      :param mode: how to randomize
 42      :param chaos: if true, the randomization-mode will be randomly changed
 43      :return: the random object instance, of the given type.
 44      """
 45      if chaos:
 46          mode = rng.choice(list(RandomizationMode))
 47      if isinstance(typ, str):
 48          # Bytes array
 49          if typ == 'bytes':
 50              if mode == RandomizationMode.mode_nil_count:
 51                  return b''
 52              if mode == RandomizationMode.mode_max_count:
 53                  return get_random_bytes_list(rng, max_bytes_length)
 54              if mode == RandomizationMode.mode_one_count:
 55                  return get_random_bytes_list(rng, 1)
 56              if mode == RandomizationMode.mode_zero:
 57                  return b'\x00'
 58              if mode == RandomizationMode.mode_max:
 59                  return b'\xff'
 60              return get_random_bytes_list(rng, rng.randint(0, max_bytes_length))
 61          elif typ[:5] == 'bytes' and len(typ) > 5:
 62              length = int(typ[5:])
 63              # Sanity, don't generate absurdly big random values
 64              # If a client is aiming to performance-test, they should create a benchmark suite.
 65              assert length <= max_bytes_length
 66              if mode == RandomizationMode.mode_zero:
 67                  return b'\x00' * length
 68              if mode == RandomizationMode.mode_max:
 69                  return b'\xff' * length
 70              return get_random_bytes_list(rng, length)
 71          # Basic types
 72          else:
 73              if mode == RandomizationMode.mode_zero:
 74                  return get_min_basic_value(typ)
 75              if mode == RandomizationMode.mode_max:
 76                  return get_max_basic_value(typ)
 77              return get_random_basic_value(rng, typ)
 78      # Vector:
 79      elif isinstance(typ, list) and len(typ) == 2:
 80          return [get_random_ssz_object(rng, typ[0], max_bytes_length, max_list_length, mode, chaos) for _ in range(typ[1])]
 81      # List:
 82      elif isinstance(typ, list) and len(typ) == 1:
 83          length = rng.randint(0, max_list_length)
 84          if mode == RandomizationMode.mode_one_count:
 85              length = 1
 86          if mode == RandomizationMode.mode_max_count:
 87              length = max_list_length
 88          return [get_random_ssz_object(rng, typ[0], max_bytes_length, max_list_length, mode, chaos) for _ in range(length)]
 89      # Container:
 90      elif hasattr(typ, 'fields'):
 91          return typ(**{field: get_random_ssz_object(rng, subtype, max_bytes_length, max_list_length, mode, chaos) for field, subtype in typ.fields.items()})
 92      else:
 93          print(typ)
 94          raise Exception("Type not recognized")
 95  
 96  
 97  def get_random_bytes_list(rng: Random, length: int) -> bytes:
 98      return bytes(rng.getrandbits(8) for _ in range(length))
 99  
100  
101  def get_random_basic_value(rng: Random, typ: str) -> Any:
102      if typ == 'bool':
103          return rng.choice((True, False))
104      if typ[:4] == 'uint':
105          size = int(typ[4:])
106          assert size in UINT_SIZES
107          return rng.randint(0, 2**size - 1)
108      if typ == 'byte':
109          return rng.randint(0, 8)
110      else:
111          raise ValueError("Not a basic type")
112  
113  
114  def get_min_basic_value(typ: str) -> Any:
115      if typ == 'bool':
116          return False
117      if typ[:4] == 'uint':
118          size = int(typ[4:])
119          assert size in UINT_SIZES
120          return 0
121      if typ == 'byte':
122          return 0x00
123      else:
124          raise ValueError("Not a basic type")
125  
126  
127  def get_max_basic_value(typ: str) -> Any:
128      if typ == 'bool':
129          return True
130      if typ[:4] == 'uint':
131          size = int(typ[4:])
132          assert size in UINT_SIZES
133          return 2**size - 1
134      if typ == 'byte':
135          return 0xff
136      else:
137          raise ValueError("Not a basic type")