CouncilEngine/test_bounty_vulnerability_engine.py at main · Simultech369/CouncilEngine · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
import unittest
from bounty_vulnerability_engine import BountyVulnerabilityEngine

class TestBountyVulnerabilityEngine(unittest.TestCase):

    def setUp(self):
        self.engine = BountyVulnerabilityEngine()

    def test_cwe_89_sql_injection_detection(self):
        vulnerable_code = """
def get_user(cursor, username):
    cursor.execute(f"SELECT * FROM users WHERE username = '{username}'")
"""
        secure_code = """
def get_user(cursor, username):
    cursor.execute("SELECT * FROM users WHERE username = %s", (username,))
"""
        res_vuln = self.engine.check_cwe_89_sql_injection(vulnerable_code)
        self.assertFalse(res_vuln.passed)
        self.assertEqual(res_vuln.cwe_id, "CWE-89")

        res_sec = self.engine.check_cwe_89_sql_injection(secure_code)
        self.assertTrue(res_sec.passed)

    def test_cwe_22_path_traversal_detection(self):
        vulnerable_code = """
def read_user_file(filename):
    with open('../data/' + filename, 'r') as f:
        return f.read()
"""
        secure_code = """
import os
def read_user_file(filename, base_dir='/safe/data'):
    canonical = os.path.realpath(os.path.join(base_dir, filename))
    if not canonical.startswith(os.path.realpath(base_dir)):
        raise ValueError('Path traversal blocked')
    with open(canonical, 'r') as f:
        return f.read()
"""
        res_vuln = self.engine.check_cwe_22_path_traversal(vulnerable_code)
        self.assertFalse(res_vuln.passed)

        res_sec = self.engine.check_cwe_22_path_traversal(secure_code)
        self.assertTrue(res_sec.passed)

    def test_cwe_502_deserialization_detection(self):
        vulnerable_code = "import pickle\ndata = pickle.loads(raw_user_bytes)"
        secure_code = "import json\ndata = json.loads(raw_user_json)"

        res_vuln = self.engine.check_cwe_502_deserialization(vulnerable_code)
        self.assertFalse(res_vuln.passed)

        res_sec = self.engine.check_cwe_502_deserialization(secure_code)
        self.assertTrue(res_sec.passed)

    def test_cwe_78_command_injection_detection(self):
        vulnerable_code = """
import subprocess
def ping_host(host):
    subprocess.run(f"ping -c 1 {host}", shell=True)
"""
        secure_code = """
import subprocess
def ping_host(host):
    subprocess.run(["/bin/ping", "-c", "1", host], shell=False)
"""
        res_vuln = self.engine.check_cwe_78_command_injection(vulnerable_code)
        self.assertFalse(res_vuln.passed)

        res_sec = self.engine.check_cwe_78_command_injection(secure_code)
        self.assertTrue(res_sec.passed)

    def test_patch_integrity_score(self):
        secure_patch = """
import json
import subprocess
def safe_func(cursor, uid):
    cursor.execute("SELECT * FROM accounts WHERE id = %s", (uid,))
    subprocess.run(["/usr/bin/id", uid], shell=False)
"""
        pis, checks = self.engine.evaluate_patch_integrity_score(secure_patch)
        self.assertEqual(pis, 1.0)
        self.assertTrue(all(c.passed for c in checks))

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