Coverage for /pythoncovmergedfiles/medio/medio/src/pdfplumber/fuzz/fuzz_helpers.py: 63%

1#!/usr/bin/python3 

2# Copyright 2023 Google LLC 

3# 

4# Licensed under the Apache License, Version 2.0 (the "License"); 

5# you may not use this file except in compliance with the License. 

6# You may obtain a copy of the License at 

7# 

8# http://www.apache.org/licenses/LICENSE-2.0 

9# 

10# Unless required by applicable law or agreed to in writing, software 

11# distributed under the License is distributed on an "AS IS" BASIS, 

12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

13# See the License for the specific language governing permissions and 

14# limitations under the License. 

15# 

16################################################################################ 

17import contextlib 

18import io 

19import tempfile 

20from enum import IntEnum 

21from typing import Protocol, Type, TypeVar 

22 

23import atheris 

24 

25 

26class HasMax(Protocol): 

27 MAX: int 

28 

29 

30T = TypeVar("T", bound=IntEnum) 

31 

32 

33class EnhancedFuzzedDataProvider(atheris.FuzzedDataProvider): 

34 def ConsumeRandomBytes(self) -> bytes: 

35 return self.ConsumeBytes(self.ConsumeIntInRange(0, self.remaining_bytes())) 

36 

37 def ConsumeRandomString(self) -> str: 

38 return self.ConsumeUnicodeNoSurrogates( 

39 self.ConsumeIntInRange(0, self.remaining_bytes()) 

40 ) 

41 

42 def ConsumeRemainingString(self) -> str: 

43 return self.ConsumeUnicodeNoSurrogates(self.remaining_bytes()) 

44 

45 def ConsumeRemainingBytes(self) -> bytes: 

46 return self.ConsumeBytes(self.remaining_bytes()) 

47 

48 @contextlib.contextmanager 

49 def ConsumeMemoryFile( 

50 self, all_data: bool = False, as_bytes: bool = True 

51 ) -> io.BytesIO: 

52 if all_data: 

53 file_data = ( 

54 self.ConsumeRemainingBytes() 

55 if as_bytes 

56 else self.ConsumeRemainingString() 

57 ) 

58 else: 

59 file_data = ( 

60 self.ConsumeRandomBytes() if as_bytes else self.ConsumeRandomString() 

61 ) 

62 

63 file = io.BytesIO(file_data) if as_bytes else io.StringIO(file_data) 

64 yield file 

65 file.close() 

66 

67 @contextlib.contextmanager 

68 def ConsumeTemporaryFile( 

69 self, suffix: str, all_data: bool = False, as_bytes: bool = True 

70 ) -> str: 

71 if all_data: 

72 file_data = ( 

73 self.ConsumeRemainingBytes() 

74 if as_bytes 

75 else self.ConsumeRemainingString() 

76 ) 

77 else: 

78 file_data = ( 

79 self.ConsumeRandomBytes() if as_bytes else self.ConsumeRandomString() 

80 ) 

81 

82 mode = "w+b" if as_bytes else "w+" 

83 tfile = tempfile.NamedTemporaryFile(mode=mode, suffix=suffix) 

84 tfile.write(file_data) 

85 tfile.seek(0) 

86 tfile.flush() 

87 yield tfile.name 

88 tfile.close() 

89 

90 def ConsumeEnum(self, enum_type: Type[T]) -> T: 

91 return enum_type(self.ConsumeIntInRange(0, enum_type.MAX))