44TODO: Leverage Hypothesis!
55"""
66from abc import ABC , abstractmethod
7+ from enum import Enum , auto
78import random
89import string
9- from typing import List , Optional , Sequence , cast
10+ from typing import List , Optional , Sequence , Type , cast
1011
1112from algosdk import abi , encoding
1213
14+
1315from graviton .models import PyTypes
1416
1517
1618class ABIStrategy (ABC ):
19+ """
20+ TODO: when incorporating hypothesis strategies, we'll need a more holistic
21+ approach that looks at relationships amongst various args.
22+ Current approach only looks at each argument as a completely independent entity.
23+ """
24+
25+ @abstractmethod
26+ def __init__ (self , abi_instance : abi .ABIType , dynamic_length : Optional [int ] = None ):
27+ pass
28+
1729 @abstractmethod
1830 def get (self ) -> PyTypes :
1931 pass
@@ -167,6 +179,15 @@ def address_logic(x):
167179
168180
169181class RandomABIStrategyHalfSized (RandomABIStrategy ):
182+ """
183+ This strategy only generates data that is half the size that _ought_ to be possible.
184+ This is useful in the case that operations involving the generated arguments
185+ could overflow due to multiplication.
186+
187+ Since this only makes sense for `abi.UintType`, it degenerates to the standard
188+ `RandomABIStrategy` for other types.
189+ """
190+
170191 def __init__ (
171192 self ,
172193 abi_instance : abi .ABIType ,
@@ -183,3 +204,229 @@ def get(self) -> PyTypes:
183204 return cast (int , full_random ) % (
184205 1 << (cast (abi .UintType , self .abi_type ).bit_size // 2 )
185206 )
207+
208+
209+ class ABIArgsMod (Enum ):
210+ # insert a random byte into selector:
211+ selector_byte_insert = auto ()
212+ # delete a byte at a random position from the selector:
213+ selector_byte_delete = auto ()
214+ # replace a random byte in the selector:
215+ selector_byte_replace = auto ()
216+ # delete a random argument:
217+ parameter_delete = auto ()
218+ # insert a random argument:
219+ parameter_append = auto ()
220+
221+
222+ class CallStrategy (ABC ):
223+ def __init__ (
224+ self ,
225+ argument_strategy : Type [ABIStrategy ] = RandomABIStrategy ,
226+ * ,
227+ num_dryruns : int = 1 ,
228+ ):
229+ self .argument_strategy : Type [ABIStrategy ] = argument_strategy
230+ self .num_dryruns : int = num_dryruns
231+
232+ def generate_value (self , gen_type : abi .ABIType ) -> PyTypes :
233+ return cast (Type [ABIStrategy ], self .argument_strategy )(gen_type ).get ()
234+
235+ @abstractmethod
236+ def generate_inputs (self , method : Optional [str ] = None ) -> List [Sequence [PyTypes ]]:
237+ pass
238+
239+
240+ class ABICallStrategy (CallStrategy ):
241+ """
242+ TODO: refactor to comport with ABIStrategy + Hypothesis
243+ TODO: make this generic on the strategy type
244+ """
245+
246+ append_args_type : abi .ABIType = abi .ByteType ()
247+
248+ def __init__ (
249+ self ,
250+ contract : str ,
251+ argument_strategy : Type [ABIStrategy ] = RandomABIStrategy ,
252+ * ,
253+ num_dryruns : int = 1 ,
254+ handle_selector : bool = True ,
255+ abi_args_mod : Optional [ABIArgsMod ] = None ,
256+ ):
257+ """
258+ contract - ABI Contract JSON
259+
260+ argument_strategy (default=RandomABIStrategy) - ABI strategy for generating arguments
261+
262+ num_dry_runs (default=1) - the number of dry runs to run
263+ (generates different inputs each time)
264+
265+ handle_selector (default=True) - usually we'll want to let
266+ `ABIContractExecutor.run_sequence()`
267+ handle adding the method selector so this param.
268+ But if set False: when providing `inputs`
269+ ensure that the 0'th argument for method calls is the selector.
270+ And when set True: when NOT providing `inputs`, the selector arg
271+ at index 0 will be added automatically.
272+
273+ abi_args_mod (optional) - when desiring to mutate the args, provide an ABIArgsMod value
274+ """
275+ super ().__init__ (argument_strategy , num_dryruns = num_dryruns )
276+ self .contract : abi .Contract = abi .Contract .from_json (contract )
277+ self .handle_selector = handle_selector
278+ self .abi_args_mod = abi_args_mod
279+
280+ def abi_method (self , method : Optional [str ]) -> abi .Method :
281+ assert method , "cannot get abi.Method for bare app call"
282+
283+ return self .contract .get_method_by_name (method )
284+
285+ def method_signature (self , method : Optional [str ]) -> Optional [str ]:
286+ """Returns None, for a bare app call (method=None signals this)"""
287+ if method is None :
288+ return None
289+
290+ return self .abi_method (method ).get_signature ()
291+
292+ def method_selector (self , method : Optional [str ]) -> bytes :
293+ assert method , "cannot get method_selector for bare app call"
294+
295+ return self .abi_method (method ).get_selector ()
296+
297+ def argument_types (self , method : Optional [str ]) -> List [abi .ABIType ]:
298+ """
299+ Argument types (excluding selector)
300+ """
301+ if method is None :
302+ return []
303+
304+ return [cast (abi .ABIType , arg .type ) for arg in self .abi_method (method ).args ]
305+
306+ def num_args (self , method : Optional [str ]) -> int :
307+ return len (self .argument_types (method ))
308+
309+ def generate_inputs (self , method : Optional [str ] = None ) -> List [Sequence [PyTypes ]]:
310+ """
311+ Generates inputs appropriate for bare app calls and method calls
312+ according to available argument_strategy.
313+ """
314+ assert (
315+ self .argument_strategy
316+ ), "cannot generate inputs without an argument_strategy"
317+
318+ mutating = self .abi_args_mod is not None
319+
320+ if not (method or mutating ):
321+ # bare calls receive no arguments (unless mutating)
322+ return [tuple () for _ in range (self .num_dryruns )]
323+
324+ arg_types = self .argument_types (method )
325+
326+ prefix : List [bytes ] = []
327+ if self .handle_selector and method :
328+ prefix = [self .method_selector (method )]
329+
330+ modify_selector = False
331+ if (action := self .abi_args_mod ) in (
332+ ABIArgsMod .selector_byte_delete ,
333+ ABIArgsMod .selector_byte_insert ,
334+ ABIArgsMod .selector_byte_replace ,
335+ ):
336+ assert (
337+ prefix
338+ ), f"{ self .abi_args_mod = } which means we need to modify the selector, but we don't have one available to modify"
339+ modify_selector = True
340+
341+ def selector_mod (prefix ):
342+ """
343+ modifies the selector by mutating a random byte (when modify_selector == True)
344+ ^^^
345+ """
346+ assert isinstance (prefix , list ) and len (prefix ) <= 1
347+ if not (prefix and modify_selector ):
348+ return prefix
349+
350+ selector = prefix [0 ]
351+ idx = random .randint (0 , 4 )
352+ x , y = selector [:idx ], selector [idx :]
353+ if action == ABIArgsMod .selector_byte_insert :
354+ selector = x + random .randbytes (1 ) + y
355+ elif action == ABIArgsMod .selector_byte_delete :
356+ selector = (x [:- 1 ] + y ) if x else y [:- 1 ]
357+ else :
358+ assert (
359+ action == ABIArgsMod .selector_byte_replace
360+ ), f"expected action={ ABIArgsMod .selector_byte_replace } but got [{ action } ]"
361+ idx = random .randint (0 , 3 )
362+ selector = (
363+ selector [:idx ]
364+ + bytes ([(selector [idx ] + 1 ) % 256 ])
365+ + selector [idx + 1 :]
366+ )
367+ return [selector ]
368+
369+ def args_mod (args ):
370+ """
371+ modifies the args by appending or deleting a random value (for appropriate `action`)
372+ ^^^
373+ """
374+ if action not in (ABIArgsMod .parameter_append , ABIArgsMod .parameter_delete ):
375+ return args
376+
377+ if action == ABIArgsMod .parameter_delete :
378+ return args if not args else tuple (args [:- 1 ])
379+
380+ assert action == ABIArgsMod .parameter_append
381+ return args + (self .generate_value (self .append_args_type ),)
382+
383+ def gen_args ():
384+ # TODO: when incorporating hypothesis strategies, we'll need a more holistic
385+ # approach that looks at relationships amongst various args
386+ args = tuple (
387+ selector_mod (prefix )
388+ + [self .generate_value (atype ) for atype in arg_types ]
389+ )
390+ return args_mod (args )
391+
392+ return [gen_args () for _ in range (self .num_dryruns )]
393+
394+
395+ class RandomArgLengthCallStrategy (CallStrategy ):
396+ """
397+ Generate a random number or arguments using the single
398+ argument_strategy provided.
399+ """
400+
401+ def __init__ (
402+ self ,
403+ argument_strategy : Type [ABIStrategy ],
404+ max_args : int ,
405+ * ,
406+ num_dryruns : int = 1 ,
407+ min_args : int = 0 ,
408+ type_for_args : abi .ABIType = abi .ABIType .from_string ("byte[8]" ),
409+ ):
410+ super ().__init__ (argument_strategy , num_dryruns = num_dryruns )
411+ self .max_args : int = max_args
412+ self .min_args : int = min_args
413+ self .type_for_args : abi .ABIType = type_for_args
414+
415+ def generate_inputs (self , method : Optional [str ] = None ) -> List [Sequence [PyTypes ]]:
416+ assert (
417+ method is None
418+ ), f"actual non-None method={ method } not supported for RandomArgLengthCallStrategy"
419+ assert (
420+ self .argument_strategy
421+ ), "cannot generate inputs without an argument_strategy"
422+
423+ def gen_args ():
424+ num_args = random .randint (self .min_args , self .max_args )
425+ abi_args = [
426+ self .generate_value (self .type_for_args ) for _ in range (num_args )
427+ ]
428+ # because cannot provide a method signature to include the arg types,
429+ # we need to convert back to raw bytes
430+ return tuple (bytes (arg ) for arg in abi_args )
431+
432+ return [gen_args () for _ in range (self .num_dryruns )]
0 commit comments