BackDoorBuster/tests/test_backdoor_detection.py

import unittest
import warnings

from detection.backdoor_detection import find_dangerous_functions
from detection.GPTdetection import detectGPT
import os


class TestBackdoorDetection(unittest.TestCase):
    def test_high_risk_detection(self):
        content = """import os
        os.system('rm -rf /')   # high risk
        exec('print("Hello")')  # high risk
        eval('2 + 2')   # high risk
        """
        file_extension = ".py"
        results = find_dangerous_functions(content, file_extension)
        self.assertIn((2, "os.system('rm -rf /')"), results["high"])
        self.assertIn((3, "exec('print(\"Hello\")')"), results["high"])
        self.assertIn((4, "eval('2 + 2')"), results["high"])

    def test_medium_risk_detection(self):
        content = """import subprocess
        subprocess.run(['ls', '-l'])    # medium risk
        import os
        os.popen('ls')  # medium risk
        """
        file_extension = ".py"
        results = find_dangerous_functions(content, file_extension)
        self.assertIn((2, "subprocess.run(['ls', '-l'])"), results["medium"])
        self.assertIn((4, "os.popen('ls')"), results["medium"])

    def test_no_risk_detection(self):
        content = """a = 10
        b = a + 5
        print('This should not be detected as risky.')
        """
        file_extension = ".py"
        results = find_dangerous_functions(content, file_extension)
        self.assertEqual(len(results["high"]), 0)
        self.assertEqual(len(results["medium"]), 0)
        self.assertEqual(len(results["low"]), 0)

    def test_inclusion_of_comments(self):
        content = """# Just a comment line
        print('This is a safe line')
        eval('2 + 2')  # This should be high risk
        subprocess.run(['echo', 'hello'])  # This should be medium risk
        """
        file_extension = ".py"
        results = find_dangerous_functions(content, file_extension)
        self.assertIn(
            (3, "eval('2 + 2')"),
            results["high"],
        )
        self.assertIn(
            (4, "subprocess.run(['echo', 'hello'])"),
            results["medium"],
        )

    def test_gpt_risk_detection(self):
        if os.getenv("OPENAI_API_KEY") is None:
            warnings.warn("OPENAI_API_KEY is not set, test skipped.", UserWarning)
            self.skipTest("OPENAI_API_KEY is not set")
        content = """import os
        os.system('rm -rf /')   # high risk
        exec('print("Hello")')  # high risk
        eval('2 + 2')   # high risk
        """
        results = detectGPT(content)
        self.assertEqual(len(results["high"]), 3)

    def test_gpt_no_risk_detection(self):
        if os.getenv("OPENAI_API_KEY") is None:
            warnings.warn("OPENAI_API_KEY is not set, test skipped.", UserWarning)
            self.skipTest("OPENAI_API_KEY is not set")
        content = """a = 10
        b = a + 5
        print('This should not be detected as risky.')
        """
        results = detectGPT(content)
        self.assertEqual(len(results["high"]), 0)
        self.assertEqual(len(results["medium"]), 0)
        self.assertEqual(len(results["low"]), 0)

    def test_gpt_env_no_set(self):
        if os.getenv("OPENAI_API_KEY") is not None:
            self.skipTest("OPENAI_API_KEY is setted")
        content = "print('test test')"
        with self.assertRaises(ValueError):
            detectGPT(content)


if __name__ == "__main__":
    unittest.main()