Spaces:
Sleeping
Sleeping
| """ | |
| Test SQL Preservation in Remediation | |
| Verifies that SQL injection fixes preserve: | |
| - Table names | |
| - Column selections | |
| - Authorization clauses (AND NOT role = 'admin') | |
| - Suffix conditions (ORDER BY, LIMIT) | |
| """ | |
| import pytest | |
| import sys | |
| import os | |
| sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) | |
| from gen_ai.model import SecureCodeGenerator | |
| from gen_ai.output_validator import CodeOutputValidator | |
| from governance.safety import SafetyChecker | |
| class TestSQLStructureExtraction: | |
| """Test the SQL structure extraction logic.""" | |
| def setup_method(self): | |
| self.generator = SecureCodeGenerator() | |
| def test_extract_table_name(self): | |
| """Verify table name is extracted correctly.""" | |
| code = "SELECT * FROM orders WHERE id = " + "user_id" | |
| info = self.generator._extract_full_sql_structure(code) | |
| assert info['table'] == 'orders' | |
| def test_extract_select_columns(self): | |
| """Verify SELECT columns are extracted.""" | |
| code = "SELECT id, email, name FROM users WHERE username = " + "input" | |
| info = self.generator._extract_full_sql_structure(code) | |
| assert 'id' in info['select_columns'] | |
| assert 'email' in info['select_columns'] | |
| def test_extract_where_column(self): | |
| """Verify WHERE column is extracted.""" | |
| code = "SELECT * FROM users WHERE user_id = " + "var" | |
| info = self.generator._extract_full_sql_structure(code) | |
| assert info['column'] == 'user_id' | |
| def test_extract_authorization_suffix(self): | |
| """Verify authorization clauses are captured.""" | |
| code = "SELECT * FROM users WHERE name = 'test' AND NOT role = 'admin'" | |
| info = self.generator._extract_full_sql_structure(code) | |
| assert "AND NOT role = 'admin'" in info['authorization_suffix'] | |
| def test_extract_order_by(self): | |
| """Verify ORDER BY is extracted.""" | |
| code = "SELECT * FROM products WHERE category = 'x' ORDER BY price DESC" | |
| info = self.generator._extract_full_sql_structure(code) | |
| assert 'ORDER BY' in info['order_by'] | |
| assert 'price' in info['order_by'] | |
| def test_extract_limit(self): | |
| """Verify LIMIT is extracted.""" | |
| code = "SELECT * FROM items WHERE active = 1 LIMIT 10" | |
| info = self.generator._extract_full_sql_structure(code) | |
| assert 'LIMIT 10' in info['limit'] | |
| class TestSQLPreservationValidation: | |
| """Test the SQL preservation validator.""" | |
| def test_table_mismatch_rejected(self): | |
| """Reject if table name changes.""" | |
| original = "SELECT * FROM orders WHERE id = " + "var" | |
| fixed = "SELECT * FROM users WHERE id = ?" | |
| is_valid, msg = CodeOutputValidator.validate_sql_preservation(original, fixed) | |
| assert not is_valid | |
| assert 'orders' in msg.lower() or 'table' in msg.lower() | |
| def test_authorization_clause_preserved(self): | |
| """Accept if authorization clause is preserved.""" | |
| original = "SELECT * FROM users WHERE name = 'x' AND NOT role = 'admin'" | |
| fixed = "SELECT * FROM users WHERE name = ? AND NOT role = 'admin'" | |
| is_valid, msg = CodeOutputValidator.validate_sql_preservation(original, fixed) | |
| assert is_valid | |
| def test_authorization_clause_lost_rejected(self): | |
| """Reject if authorization clause is lost.""" | |
| original = "SELECT * FROM users WHERE name = 'x' AND NOT role = 'admin'" | |
| fixed = "SELECT * FROM users WHERE name = ?" | |
| is_valid, msg = CodeOutputValidator.validate_sql_preservation(original, fixed) | |
| assert not is_valid | |
| assert 'authorization' in msg.lower() or 'AND NOT' in msg | |
| def test_order_by_preserved(self): | |
| """Accept if ORDER BY is preserved.""" | |
| original = "SELECT * FROM products WHERE cat = 'x' ORDER BY price" | |
| fixed = "SELECT * FROM products WHERE cat = ? ORDER BY price" | |
| is_valid, msg = CodeOutputValidator.validate_sql_preservation(original, fixed) | |
| assert is_valid | |
| def test_order_by_lost_rejected(self): | |
| """Reject if ORDER BY is lost.""" | |
| original = "SELECT * FROM products WHERE cat = 'x' ORDER BY price" | |
| fixed = "SELECT * FROM products WHERE cat = ?" | |
| is_valid, msg = CodeOutputValidator.validate_sql_preservation(original, fixed) | |
| assert not is_valid | |
| assert 'ORDER BY' in msg | |
| def test_limit_preserved(self): | |
| """Accept if LIMIT is preserved.""" | |
| original = "SELECT * FROM items WHERE x = 'y' LIMIT 10" | |
| fixed = "SELECT * FROM items WHERE x = ? LIMIT 10" | |
| is_valid, msg = CodeOutputValidator.validate_sql_preservation(original, fixed) | |
| assert is_valid | |
| class TestGovernanceEnforcement: | |
| """Test governance rejects non-compliant fixes.""" | |
| def setup_method(self): | |
| self.checker = SafetyChecker() | |
| def test_governance_rejects_table_change(self): | |
| """Governance should reject if table name changes.""" | |
| original = "SELECT * FROM orders WHERE id = " + "var" | |
| fixed = "SELECT * FROM users WHERE id = ?" | |
| is_approved, msg = self.checker.check_structure_preservation(original, fixed, "CWE-89") | |
| assert not is_approved | |
| assert 'GOVERNANCE' in msg.upper() or 'table' in msg.lower() | |
| def test_governance_approves_valid_fix(self): | |
| """Governance should approve correct parameterization.""" | |
| original = "SELECT * FROM users WHERE name = 'x'" | |
| fixed = "SELECT * FROM users WHERE name = ?" | |
| is_approved, msg = self.checker.check_structure_preservation(original, fixed, "CWE-89") | |
| assert is_approved | |
| class TestRemediationPreservation: | |
| """Test end-to-end remediation preserves structure.""" | |
| def setup_method(self): | |
| self.generator = SecureCodeGenerator() | |
| def test_python_fix_preserves_table(self): | |
| """Python fix should preserve table name.""" | |
| code = "cursor.execute('SELECT * FROM orders WHERE user_id = ' + uid)" | |
| result = self.generator.generate_fix_with_fallback(code, "CWE-89", "python") | |
| assert 'orders' in result['code'] | |
| def test_python_fix_preserves_columns(self): | |
| """Python fix should preserve column selections.""" | |
| code = "cursor.execute('SELECT id, email FROM users WHERE name = ' + name)" | |
| result = self.generator.generate_fix_with_fallback(code, "CWE-89", "python") | |
| # Should have parameterization | |
| assert '?' in result['code'] | |
| # Should preserve table | |
| assert 'users' in result['code'] | |
| def test_java_fix_preserves_table(self): | |
| """Java fix should preserve table name.""" | |
| code = 'stmt.executeQuery("SELECT * FROM products WHERE id = " + productId)' | |
| result = self.generator.generate_fix_with_fallback(code, "CWE-89", "java") | |
| assert 'products' in result['code'] | |
| assert 'PreparedStatement' in result['code'] | |
| if __name__ == "__main__": | |
| pytest.main([__file__, "-v"]) | |