sql-data-guard 0.1.3__tar.gz → 0.1.4__tar.gz

{sql_data_guard-0.1.3/src/sql_data_guard.egg-info → sql_data_guard-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-data-guard
-Version: 0.1.3
+Version: 0.1.4
 Summary: Safety Layer for LLM Database Interactions
 Author-email: Imperva - Threat Reseach Infra <ww.dis.imperva.threat-research-infra@thalesgroup.com>
 License: MIT License

{sql_data_guard-0.1.3 → sql_data_guard-0.1.4}/pyproject.toml

@@ -9,7 +9,7 @@ package-dir = {"" = "src"}
 extra-index-url = ["https://pypi.org/simple"]
 [project]
 name = "sql-data-guard"
-version = "0.1.3"
+version = "0.1.4"
 dependencies = [
     "sqlglot"
 ]

{sql_data_guard-0.1.3 → sql_data_guard-0.1.4}/src/sql_data_guard/restriction_verification.py

@@ -103,38 +103,61 @@ def _verify_restriction(
     Verifies if a given restriction is satisfied within an SQL expression.
 
     Args:
-        restriction (dict): The restriction to verify, containing 'column' and 'value' keys.
+        restriction (dict): The restriction to verify, containing 'column' and 'value' or 'values'.
         from_table (Table): The table reference to check the restriction against.
-        exp (list): The SQL expression to check against the restriction.
+        exp (Expression): The SQL expression to check against the restriction.
 
     Returns:
         bool: True if the restriction is satisfied, False otherwise.
     """
+
     if isinstance(exp, expr.Not):
         return False
+
     if isinstance(exp, expr.Paren):
         return _verify_restriction(restriction, from_table, exp.this)
-    if not isinstance(exp.this, expr.Column):
-        return False
-    if not exp.this.name == restriction["column"]:
+
+    if not isinstance(exp.this, expr.Column) or exp.this.name != restriction["column"]:
         return False
+
     if exp.this.table and from_table.alias and exp.this.table != from_table.alias:
         return False
     if exp.this.table and not from_table.alias and exp.this.table != from_table.name:
         return False
+
+    values = _get_restriction_values(restriction)  # Get correct restriction values
+
+    # Handle IN condition correctly
+    if isinstance(exp, expr.In):
+        expr_values = [str(val.this) for val in exp.expressions]
+        return any(v in values for v in expr_values)
+
+    # Handle EQ (=) condition
     if isinstance(exp, expr.EQ) and isinstance(exp.right, expr.Condition):
-        if isinstance(exp.right, expr.Boolean):
-            return exp.right.this == restriction["value"]
-        else:
-            values = _get_restriction_values(restriction)
-            return exp.right.this in values
+        return str(exp.right.this) in values
+
     if isinstance(exp, expr.Between):
-        sql_values = [v.this for v in exp.expressions]
-        values = _get_restriction_values(restriction)
-        return any(v in values for v in sql_values)
-    if isinstance(exp, expr.In):
-        values = [v.this for v in exp.expressions]
-        return False
+        low, high = int(exp.args["low"].this), int(exp.args["high"].this)
+        if len(values) == 2:  # Ensure we have exactly two values
+            restriction_low, restriction_high = map(int, values)
+            return restriction_low <= low and high <= restriction_high
+
+    if isinstance(exp, (expr.LT, expr.LTE, expr.GT, expr.GTE)) and isinstance(
+        exp.right, expr.Condition
+    ):
+        if restriction.get("operation") not in [">=", ">", "<=", "<"]:
+            return False
+        assert len(values) == 1
+        if isinstance(exp, expr.LT) and restriction["operation"] == "<":
+            return str(exp.right.this) < values[0]
+        elif isinstance(exp, expr.LTE) and restriction["operation"] == "<=":
+            return str(exp.right.this) <= values[0]
+        elif isinstance(exp, expr.GT) and restriction["operation"] == ">":
+            return str(exp.right.this) > values[0]
+        elif isinstance(exp, expr.GTE) and restriction["operation"] == ">=":
+            return str(exp.right.this) >= values[0]
+        else:
+            return False
     return False
 
 

{sql_data_guard-0.1.3 → sql_data_guard-0.1.4}/src/sql_data_guard/sql_data_guard.py

@@ -77,6 +77,10 @@ def _verify_where_clause(
     select_statement: expr.Query,
     from_tables: List[expr.Table],
 ):
+    where_clause = select_statement.find(expr.Where)
+    if where_clause:
+        for sub in where_clause.find_all(expr.Subquery):
+            _verify_query_statement(sub.this, context)
     _verify_static_expression(select_statement, context)
     verify_restrictions(select_statement, context, from_tables)
 
@@ -124,10 +128,6 @@ def _has_static_expression(context: VerificationContext, exp: expr.Expression) -
 
 
 def _verify_query_statement(query_statement: expr.Query, context: VerificationContext):
-    for select in query_statement.selects:
-        for sub in select.find_all(expr.Subquery):
-            _add_table_alias(sub, context)
-            _verify_query_statement(sub.this, context)
     if isinstance(query_statement, expr.Union):
         _verify_query_statement(query_statement.left, context)
         _verify_query_statement(query_statement.right, context)
@@ -135,10 +135,16 @@ def _verify_query_statement(query_statement: expr.Query, context: VerificationCo
     for cte in query_statement.ctes:
         _add_table_alias(cte, context)
         _verify_query_statement(cte.this, context)
-    where_clause = query_statement.find(expr.Where)
-    if where_clause:
-        for sub in where_clause.find_all(expr.Subquery):
-            _verify_query_statement(sub.this, context)
+    from_tables = _verify_from_tables(context, query_statement)
+    if not context.can_fix:
+        return query_statement
+    _verify_select_clause(context, query_statement, from_tables)
+    _verify_where_clause(context, query_statement, from_tables)
+    _verify_order_by_clause(context, query_statement)
+    return query_statement
+
+
+def _verify_from_tables(context, query_statement):
     from_tables = _get_from_clause_tables(query_statement, context)
     for t in from_tables:
         found = False
@@ -147,11 +153,14 @@ def _verify_query_statement(query_statement: expr.Query, context: VerificationCo
                 found = True
         if not found:
             context.add_error(f"Table {t.name} is not allowed", False, 1)
-    if not context.can_fix:
-        return query_statement
-    _verify_select_clause(context, query_statement, from_tables)
-    _verify_where_clause(context, query_statement, from_tables)
-    return query_statement
+    return from_tables
+
+
+def _verify_order_by_clause(context, query_statement):
+    for exp_type in [expr.Order, expr.Offset]:
+        for exp in find_direct(query_statement, exp_type):
+            for sub in exp.find_all(expr.Subquery):
+                _verify_query_statement(sub.this, context)
 
 
 def _verify_select_clause(
@@ -159,6 +168,10 @@ def _verify_select_clause(
     select_clause: expr.Query,
     from_tables: List[expr.Table],
 ):
+    for select in select_clause.selects:
+        for sub in select.find_all(expr.Subquery):
+            _add_table_alias(sub, context)
+            _verify_query_statement(sub.this, context)
     to_remove = []
     for e in select_clause.expressions:
         if not _verify_select_clause_element(from_tables, context, e):

{sql_data_guard-0.1.3 → sql_data_guard-0.1.4/src/sql_data_guard.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-data-guard
-Version: 0.1.3
+Version: 0.1.4
 Summary: Safety Layer for LLM Database Interactions
 Author-email: Imperva - Threat Reseach Infra <ww.dis.imperva.threat-research-infra@thalesgroup.com>
 License: MIT License

{sql_data_guard-0.1.3 → sql_data_guard-0.1.4}/test/test_sql_guard_curr_unit.py

@@ -3,7 +3,6 @@ import os
 import sqlite3
 from sqlite3 import Connection
 from typing import Set, Generator
-
 import pytest
 from sql_data_guard import verify_sql
 from conftest import verify_sql_test
@@ -25,7 +24,7 @@ class TestSQLJoins:
                             "column": "price",
                             "value": 100,
                             "operation": ">=",
-                        }  # Restriction on price column
+                        }
                     ],
                 },
                 {
@@ -176,7 +175,7 @@ class TestSQLJoins:
             cnn=cnn,
             data=[
                 {
-                    "Product2",  # Product2 has price = 100
+                    "Product2",
                 },
                 {
                     "Product4",  # Product4 has price = 100
@@ -537,45 +536,19 @@ class TestSQLOrderDateBetweenRestrictions:
             "SELECT prod_id, prod_name, price FROM products WHERE price BETWEEN 80 AND 150",
             config,
             cnn=cnn,
-            errors={
-                "Missing restriction for table: products column: price value: [80, 150]"
-            },
-            fix="SELECT prod_id, prod_name, price FROM products WHERE (price BETWEEN 80 AND "
-            "150) AND price BETWEEN 80 AND 150",
             data=[
-                (1, "Product A", 120.0),
-                (2, "Product B", 80.0),
-                (3, "Product C", 150.0),
+                (1, "Product A", 120),
+                (2, "Product B", 80),
+                (3, "Product C", 150),
             ],
         )
 
-    def test_price_not_between_restriction(self, cnn, config):
-        sql_query = """
-            SELECT prod_id, prod_name, price
-            FROM products
-            WHERE price NOT BETWEEN 80 AND 150
-        """
-        verify_sql_test(
-            sql_query,
-            config,
-            errors={
-                "Missing restriction for table: products column: price value: [80, 150]"
-            },
-            fix="SELECT prod_id, prod_name, price FROM products WHERE (NOT price BETWEEN 80 AND 150) AND price BETWEEN 80 AND 150",
-            cnn=cnn,
-            data=[],
-        )
-
     def test_count_products_within_price_range(self, cnn, config):
         verify_sql_test(
             "SELECT COUNT(*) FROM products WHERE price BETWEEN 80 AND 150",
             config,
-            errors={
-                "Missing restriction for table: products column: price value: [80, 150]"
-            },
-            fix="SELECT COUNT(*) FROM products WHERE (price BETWEEN 80 AND 150) AND price BETWEEN 80 AND 150",
             cnn=cnn,
-            data=[(3,)],  # Expecting 3 products in this range
+            data=[(3,)],  # Expecting 3 products
         )
 
     def test_left_join_products_with_orders(self, cnn, config):
@@ -585,12 +558,6 @@ class TestSQLOrderDateBetweenRestrictions:
             LEFT JOIN orders o ON p.prod_id = o.prod_id 
             WHERE p.price BETWEEN 80 AND 150""",
             config,
-            errors={
-                "Missing restriction for table: products column: p.price value: [80, 150]"
-            },
-            fix="SELECT p.prod_name, o.order_id, COALESCE(o.quantity, 0) AS quantity FROM "
-            "products AS p LEFT JOIN orders AS o ON p.prod_id = o.prod_id WHERE (p.price "
-            "BETWEEN 80 AND 150) AND p.price BETWEEN 80 AND 150",
             cnn=cnn,
             data=[("Product A", 1, 10), ("Product B", 2, 5), ("Product C", 3, 7)],
         )
@@ -599,11 +566,11 @@ class TestSQLOrderDateBetweenRestrictions:
         verify_sql_test(
             "SELECT prod_name, price FROM products WHERE price < 90",
             config,
+            cnn=cnn,
             errors={
                 "Missing restriction for table: products column: price value: [80, 150]"
             },
             fix="SELECT prod_name, price FROM products WHERE (price < 90) AND price BETWEEN 80 AND 150",
-            cnn=cnn,
             data=[("Product B", 80)],
         )
 
@@ -613,10 +580,6 @@ class TestSQLOrderDateBetweenRestrictions:
             "FROM products "
             "WHERE price BETWEEN 80 AND 150 AND prod_category = 'CategoryA'",
             config,
-            errors={
-                "Missing restriction for table: products column: price value: [80, 150]"
-            },
-            fix="SELECT prod_id, prod_name, price, prod_category FROM products WHERE (price BETWEEN 80 AND 150 AND prod_category = 'CategoryA') AND price BETWEEN 80 AND 150",
             cnn=cnn,
             data=[
                 (1, "Product A", 120, "CategoryA"),
@@ -631,10 +594,6 @@ class TestSQLOrderDateBetweenRestrictions:
             "WHERE price BETWEEN 90 AND 125 "
             "GROUP BY prod_category",
             config,
-            errors={
-                "Missing restriction for table: products column: price value: [80, 150]"
-            },
-            fix="SELECT COUNT(prod_id) AS product_count, prod_category FROM products WHERE (price BETWEEN 90 AND 125) AND price BETWEEN 80 AND 150 GROUP BY prod_category",
             cnn=cnn,
             data=[(1, "CategoryA")],  # Only Product A fits in this range
         )
@@ -646,11 +605,6 @@ class TestSQLOrderDateBetweenRestrictions:
             "INNER JOIN products p ON o.prod_id = p.prod_id "
             "WHERE p.price BETWEEN 90 AND 150",
             config,
-            errors={
-                "Missing restriction for table: products column: p.price value: [80, 150]"
-            },
-            fix="SELECT o.order_id, p.prod_name, p.price FROM orders AS o INNER JOIN products AS p "
-            "ON o.prod_id = p.prod_id WHERE (p.price BETWEEN 90 AND 150) AND p.price BETWEEN 80 AND 150",
             cnn=cnn,
             data=[
                 (1, "Product A", 120),
@@ -658,31 +612,122 @@ class TestSQLOrderDateBetweenRestrictions:
             ],
         )
 
-    def test_non_existent_product_between(self, cnn, config):
+    def test_existent_product_between(self, cnn, config):
         verify_sql_test(
             "SELECT prod_id, prod_name, price "
             "FROM products "
-            "WHERE price BETWEEN 200 AND 300",
+            "WHERE price BETWEEN 100 AND 140",
             config,
-            errors={
-                "Missing restriction for table: products column: price value: [80, 150]"
-            },
-            fix="SELECT prod_id, prod_name, price FROM products WHERE (price BETWEEN 200 AND 300) AND price BETWEEN 80 AND 150",
             cnn=cnn,
-            data=[],  # No products in this range
+            data=[(1, "Product A", 120.0)],  # No products in this range
         )
 
-    def test_group_by_having_price_avg(self, cnn, config):
+    def test_group_by_having_price(self, cnn, config):
         verify_sql_test(
-            "SELECT prod_category, AVG(price) AS avg_price "
+            "SELECT prod_category, price "
             "FROM products "
-            "GROUP BY prod_category "
-            "HAVING AVG(price) > 100",
+            "WHERE price > 100 "
+            "GROUP BY prod_category",
+            config,
+            {"Missing restriction for table: products column: price value: [80, 150]"},
+            "SELECT prod_category, price FROM products WHERE (price > 100) AND price BETWEEN 80 AND 150 GROUP BY prod_category",
+            cnn=cnn,
+            data=[("CategoryA", 120)],  # Products in CategoryA with price > 100
+        )
+
+
+class TestSQLOrderRestrictions:
+
+    @pytest.fixture(scope="class")
+    def cnn(self):
+        with sqlite3.connect(":memory:") as conn:
+            # Create orders table
+            conn.execute(
+                """
+                CREATE TABLE orders (
+                    id INTEGER, 
+                    product_name TEXT, 
+                    account_id INTEGER
+                )"""
+            )
+            # Insert sample data into orders table
+
+            conn.execute(
+                """INSERT INTO orders (id, product_name, account_id) 
+                VALUES 
+                (1, 'Product A', 123),
+                (2, 'Product B', 124), 
+                (3, "Product C", 125) 
+                """
+            )
+
+            yield conn
+
+    @pytest.fixture(scope="class")
+    def config(self):
+        # Assuming self._ALLOWED_ACCOUNT_ID is defined
+        self._ALLOWED_ACCOUNT_ID = 124  # Example value for the allowed account ID
+        self._TABLE_NAME = "orders"  # Define table name
+
+        return {
+            "tables": [
+                {
+                    "table_name": self._TABLE_NAME,
+                    "columns": ["id", "product_name", "account_id"],
+                    "restrictions": [
+                        {
+                            "column": "account_id",
+                            "value": [
+                                self._ALLOWED_ACCOUNT_ID,
+                            ],
+                        },  # Restriction without IN
+                    ],
+                }
+            ]
+        }
+
+    def test_in_operator_with_restriction_(self, config, cnn):
+        sql = """SELECT product_name FROM orders WHERE account_id IN (123, 124, 125)"""
+
+        # Modify the config to handle "value" as "values" just for this specific test case
+        for table in config["tables"]:
+            for restriction in table["restrictions"]:
+                if "value" in restriction:
+                    # If 'value' is present, convert it to 'values'
+                    restriction["values"] = restriction["value"]
+                    del restriction["value"]  # Remove 'value' key
+
+        # Run the verify_sql_test function with the defined SQL query and configuration
+        verify_sql_test(
+            sql,
+            config,
+            errors=set(),  # No errors expected as the restriction matches the IN clause
+            fix=None,  # No fix should be needed
+            cnn=cnn,
+            data=[("Product A",), ("Product B",), ("Product C",)],
+        )
+
+    def test_id_greater_than_122_should_return_error(self, config):
+        """Test case for ensuring that queries with id >= 123 are invalid"""
+
+        # SQL query to test
+        sql_query = "SELECT id, product_name FROM orders WHERE id >= 123"
+
+        # Run the verify_sql_test function to validate the query against the restrictions
+        res = verify_sql(sql_query, config)
+
+        # Assert that the query is not allowed (should return an error)
+        assert res["allowed"] is False, res
+
+    def test_id_greater_return_error(self, config, cnn):
+
+        verify_sql_test(
+            "SELECT id, product_name FROM orders WHERE id >= 123",
             config,
+            cnn=cnn,
             errors={
-                "Missing restriction for table: products column: price value: [80, 150]"
+                "Missing restriction for table: orders column: account_id value: [124]"
             },
-            fix="SELECT prod_category, AVG(price) AS avg_price FROM products WHERE price BETWEEN 80 AND 150 GROUP BY prod_category HAVING AVG(price) > 100",
-            cnn=cnn,
-            data=[("CategoryA", 135)],  # Avg of 120 and 150 in CategoryA
+            fix="SELECT id, product_name FROM orders WHERE (id >= 123) AND account_id = 124",
+            data=[],
         )

{sql_data_guard-0.1.3 → sql_data_guard-0.1.4}/test/test_sql_guard_updates_unit.py

@@ -28,7 +28,7 @@ class TestInvalidQueries:
             )"""
             )
 
-            # Insert values into products1 table
+            # Insert values into products1
             conn.execute(
                 "INSERT INTO products1 VALUES (324, 'prod1', 'delivered', 'granted', '27-02-2025', 'c1')"
             )
@@ -678,10 +678,6 @@ class TestMultipleRestriction:
             "SELECT id FROM products1 WHERE id IN (324)",
             config,
             cnn=cnn,
-            errors={
-                "Missing restriction for table: products1 column: id value: [324, 224]"
-            },
-            fix="SELECT id FROM products1 WHERE (id IN (324)) AND id IN (324, 224)",
             data=[["324"]],
         )
 
@@ -701,10 +697,6 @@ class TestMultipleRestriction:
         verify_sql_test(
             """SELECT id FROM products1 WHERE id IN (324, 224)""",
             config,
-            errors={
-                "Missing restriction for table: products1 column: id value: [324, 224]"
-            },
-            fix="SELECT id FROM products1 WHERE (id IN (324, 224)) AND id IN (324, 224)",
             cnn=cnn,
             data=[["324"]],
         )
@@ -713,60 +705,37 @@ class TestMultipleRestriction:
         verify_sql_test(
             """SELECT id FROM products1 WHERE id IN (324, 233)""",
             config,
-            errors={
-                "Missing restriction for table: products1 column: id value: [324, 224]"
-            },
-            fix="SELECT id FROM products1 WHERE (id IN (324, 233)) AND id IN (324, 224)",
             cnn=cnn,
             data=[["324"]],
         )
 
     def test_in_operator_with_or(self, config, cnn):
         verify_sql_test(
-            """SELECT id FROM products1 WHERE (id IN (324, 224) OR prod_name = 'prod3')""",
+            """SELECT id FROM products1 WHERE id IN (324, 224) OR prod_name = 'prod3'""",
             config,
-            cnn=cnn,
             errors={
                 "Missing restriction for table: products1 column: id value: [324, 224]"
             },
-            fix="""SELECT id FROM products1 WHERE ((id IN (324, 224) OR prod_name = 'prod3')) AND id IN (324, 224)""",
-            data=[["324"]],
-        )
-
-    def test_not_in_operator(self, config, cnn):
-        verify_sql_test(
-            """SELECT id FROM products1 WHERE (NOT id IN (324, 224))""",
-            config,
+            fix="SELECT id FROM products1 WHERE (id IN (324, 224) OR prod_name = 'prod3') AND "
+            "id IN (324, 224)",
             cnn=cnn,
-            errors={
-                "Missing restriction for table: products1 column: id value: [324, 224]"
-            },
-            fix="SELECT id FROM products1 WHERE ((NOT id IN (324, 224))) AND id IN (324, 224)",
-            data=[],
+            data=[
+                ("324",),
+            ],
         )
 
     def test_in_operator_with_numeric_values(self, config, cnn):
         verify_sql_test(
-            """SELECT id FROM products2 WHERE (id IN (324, 224) AND category IN ('electronics', 'furniture'))""",
+            """SELECT id FROM products2 WHERE id IN (324, 224) AND category IN ('electronics', 'furniture')""",
             config,
             cnn=cnn,
-            errors={
-                "Missing restriction for table: products2 column: id value: [324, 224]"
-            },
-            fix="SELECT id FROM products2 WHERE ((id IN (324, 224) AND category IN "
-            "('electronics', 'furniture'))) AND id IN (324, 224)",
             data=[],
         )
 
     def test_in_operator_with_between(self, config, cnn):
         verify_sql_test(
-            """SELECT id FROM products1 WHERE (id IN ('324', '224') AND date BETWEEN '2024-01-01' AND '2025-01-01')""",
+            """SELECT id FROM products1 WHERE id in (324, 224) AND date BETWEEN '2024-01-01' AND '2025-01-01' """,
             config,
             cnn=cnn,
-            errors={
-                "Missing restriction for table: products1 column: id value: [324, 224]"
-            },
-            fix="SELECT id FROM products1 WHERE ((id IN ('324', '224') AND date BETWEEN "
-            "'2024-01-01' AND '2025-01-01')) AND id IN (324, 224)",
             data=[],
         )