From 5a228e5cb0adf5dc71619adc31d705e5f6b7d140 Mon Sep 17 00:00:00 2001
From: sangge-redmi <2251250136@qq.com>
Date: Tue, 4 Jun 2024 11:34:43 +0800
Subject: [PATCH] feat: update return content

---
 tests/final_tests_util.py | 61 +++++++++++++++++++++++----------------
 tests/test_final_tests.py |  3 +-
 2 files changed, 38 insertions(+), 26 deletions(-)

diff --git a/tests/final_tests_util.py b/tests/final_tests_util.py
index 62ffdcc..adbd9e6 100644
--- a/tests/final_tests_util.py
+++ b/tests/final_tests_util.py
@@ -1,7 +1,11 @@
-from typing import Tuple
+from typing import Tuple, List
 from git import Repo  # type: ignore
 import random
 from pathlib import Path
+import pickle
+import marshal
+import importlib.util
+import os
 
 
 def clone_repo(repo_url: str, clone_dir: str) -> None:
@@ -18,9 +22,12 @@ def clone_repo(repo_url: str, clone_dir: str) -> None:
         print(f"Error cloning repository: {e}")
 
 
+# a return type of backdoor. Include injected file name and number.
+
+
 def inject_random_backdoor(
     path: str, pickle: bool = False, pyc: bool = False, sample_rate: float = 0.1
-) -> int:
+) -> Tuple[Tuple[str, int], ...]:
     """
     Insert random backdoor into the path.
 
@@ -43,21 +50,22 @@ def inject_random_backdoor(
         if random.random() < sample_rate:
             injected_python_files.append(python_file)
     injectedNum = len(injected_python_files)
+
+    results: List[Tuple[str, int]] = []
+    # inject backdoor
     for python_file in injected_python_files:
         with open(python_file, "r+", errors="ignore") as f:
             lines = f.readlines()
             total_lines = len(lines)
             inject_line_number = random.randint(0, total_lines)
+            # choose random backdoor
             inject_code = random.choice(backdoors)
-
             lines.insert(inject_line_number, inject_code + "\n")
-
-            # Move the file pointer to the beginning of the file
             f.seek(0)
-
-            # Write the modified content back to the file
             f.writelines(lines)
-    return injectedNum
+            results.append((str(python_file), inject_line_number))
+
+    return tuple(results)
 
 
 # 示例1: 通过 os.system 执行命令
@@ -144,7 +152,10 @@ def inject_pickle_backdoor(path: str) -> None:
     Args:
         path (str): The path to the repository to insert the backdoor into.
     """
-    pass
+    for i, backdoor in enumerate(backdoors):
+        filename = os.path.join(path, f"backdoor{i}.pickle")
+        with open(filename, "wb") as f:
+            pickle.dump(backdoor, f)
 
 
 def inject_pyc_backdoor(path: str) -> None:
@@ -154,25 +165,25 @@ def inject_pyc_backdoor(path: str) -> None:
     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")
+
+        # Compile the string to a code object
+        code = compile(backdoor, filename, "exec")
+
+        # Create a code object header
+        header = importlib.util.MAGIC_NUMBER
+        if hasattr(importlib.util, "SOURCE_SUFFIXES"):
+            header += b"\x00" * 4
+
+        # Write the .pyc file
+        with open(filename, "wb") as file:
+            file.write(header)
+            marshal.dump(code, file)
+
     pass
 
 
-def check_accuracy(report_file: str, backdoor_location: Tuple[str, int]) -> float:
-    """
-    Check the accuracy of the backdoor insertion.
-
-    Args:
-        report_file (str): The path to the report file.
-        backdoor_location (Tuple[str, int]): The location of the backdoor in the repository.
-
-    Returns:
-        float: The accuracy rate of the backdoor insertion.
-    """
-    accuracy_rate = 0.0
-
-    return accuracy_rate
-
-
 if __name__ == "__main__":
     repo_url = "https://github.com/TheAlgorithms/Python.git"
     clone_dir = "/tmp/repo"
diff --git a/tests/test_final_tests.py b/tests/test_final_tests.py
index 3be4175..53075dd 100644
--- a/tests/test_final_tests.py
+++ b/tests/test_final_tests.py
@@ -12,7 +12,8 @@ class TestFinalTests(unittest.TestCase):
         shutil.rmtree("./tmp/repo", ignore_errors=True)
         clone_repo("https://github.com/injetlee/Python.git", "./tmp/repo")
         sampleRate = 0.1
-        self.injectedNum = inject_random_backdoor("./tmp/repo", sample_rate=sampleRate)
+        self.inject_reslt = inject_random_backdoor("./tmp/repo", sample_rate=sampleRate)
+        self.injectedNum = len(self.inject_reslt)
         print(self.injectedNum)
         project_path = Path("./tmp/repo")
         self.all_python_files = list(project_path.rglob("*.py"))