sangge-redmi
ebfc70eeae
Some checks failed
Python application test / build (pull_request) Failing after 12m13s
96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
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()
|