Coverage for src / renaissance / common / rewriter.py: 72%

1import sys 

2 

3 

4class Rewrite: 

5 def __init__(self, start: int, end: int, replacement: bytes) -> None: 

6 self.start = start 

7 self.end = end 

8 self.replacement = replacement 

9 

10 

11class Rewriter: 

12 """A class that allows for modifications to a byte sequence.""" 

13 

14 def __init__(self, content: bytes) -> None: 

15 self.__content = content 

16 self.__rewrites: list[Rewrite] = [] 

17 

18 def replace(self, start: int, end: int, new_content: bytes) -> None: 

19 """Replaces a portion of the content with new content. 

20 

21 This method will replace the content between the specified start and end 

22 indices with the provided new_content. If there is an existing rewrite 

23 that partially overlaps with the specified range, the new content will be 

24 appended to the existing replacement, and the range will be adjusted to 

25 encompass both the old and new content. If the start or end indices are out 

26 of bounds, then new content will be inserted at the end of the byte sequence. 

27 

28 Args: 

29 start (int): The starting index of the content to be replaced. 

30 end (int): The ending index of the content to be replaced. 

31 new_content (bytes): The new content to insert in place of the old content. 

32 

33 Returns: 

34 None 

35 

36 """ 

37 for r in self.__rewrites: 

38 # if r partially overlaps with start and end then append the new content to the existing replacement 

39 if r.start <= start <= r.end: 

40 r.replacement += new_content 

41 r.start = min(r.start, start) 

42 r.end = max(r.end, end) 

43 return 

44 real_start = len(self.__content) if start > len(self.__content) or start < 0 else start 

45 real_end = len(self.__content) if end > len(self.__content) or end < 0 else end 

46 self.__rewrites.append(Rewrite(real_start, real_end, new_content)) 

47 

48 def apply(self) -> bytes: 

49 """Applies the rewrites to a copied byte sequence. 

50 

51 This method reverses the order of the rewrites to ensure that insertions 

52 are performed correctly. It then sorts the rewrites by their start position 

53 in descending order and applies each rewrite to the byte sequence. 

54 

55 Returns: 

56 bytes: The modified byte sequence after all rewrites have been applied. 

57 

58 """ 

59 result = bytearray(self.__content[:]) 

60 for rewrite in sorted(self.__rewrites, key=lambda x: x.start, reverse=True): 

61 result[rewrite.start : rewrite.end] = rewrite.replacement 

62 return result 

63 

64 @property 

65 def content(self) -> bytes: 

66 return self.__content 

67 

68 

69if __name__ == "__main__": 

70 # create a byte array a random bytes of len 20 

71 

72 my_bytes = bytearray(20) 

73 for i in range(20): 

74 my_bytes[i] = ord("a") + i 

75 rewriter = Rewriter(my_bytes) 

76 rewriter.replace(5, 10, b"hellooo") 

77 rewriter.replace(5, 10, b" world") 

78 rewriter.replace(0, 0, b"BEGIN") 

79 s = rewriter.apply().decode(sys.getfilesystemencoding()) 

80 print(len(s)) 

81 print(s)