sql-data-guard 0.1.7__tar.gz → 0.1.9__tar.gz

{sql_data_guard-0.1.7/src/sql_data_guard.egg-info → sql_data_guard-0.1.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-data-guard
-Version: 0.1.7
+Version: 0.1.9
 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.7 → sql_data_guard-0.1.9}/pyproject.toml

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

sql_data_guard-0.1.9/src/sql_data_guard/mcpwrapper/mcp_wrapper.py

@@ -0,0 +1,129 @@
+import json
+import os
+import sys
+import threading
+from typing import Optional, Dict, List
+
+import docker
+
+from sql_data_guard import verify_sql
+
+
+def load_config() -> dict:
+    return json.load(open("/conf/config.json"))
+
+
+def _get_volumes() -> List[str]:
+    volumes = config["mcp-server"].get("volumes", [])
+    if "PWD" in os.environ:
+        volumes = [v.replace("$PWD", os.environ["PWD"]) for v in volumes]
+    return volumes
+
+
+def start_inner_container():
+    client = docker.from_env()
+    container = client.containers.run(
+        config["mcp-server"]["image"],
+        (
+            " ".join(config["mcp-server"]["args"])
+            if "args" in config["mcp-server"]
+            else None
+        ),
+        volumes=_get_volumes(),
+        network_mode=config["mcp-server"].get("network-mode"),
+        stdin_open=True,
+        auto_remove=True,
+        detach=True,
+        stdout=True,
+    )
+
+    def stream_output_inject_response():
+        for line in container.logs(stream=True):
+            line_json = json.loads(line)
+            request_id = line_json["id"]
+            if request_id in errors:
+                if "result" in line_json and "content" in line_json["result"]:
+                    line_json["result"]["content"].insert(
+                        0,
+                        {
+                            "type": "text",
+                            "text": f"[{errors[request_id]}]",
+                            "isError": True,
+                        },
+                    )
+                del errors[request_id]
+            sys.stdout.write(json.dumps(line_json) + "\n")
+            sys.stdout.flush()
+
+    def stream_output():
+        for line in container.logs(stream=True):
+            sys.stdout.write(line.decode("utf-8"))
+            sys.stdout.flush()
+
+    threading.Thread(
+        target=stream_output_inject_response if inject_response else stream_output,
+        daemon=True,
+    ).start()
+    return container
+
+
+def main():
+    container = start_inner_container()
+    try:
+        socket = container.attach_socket(params={"stdin": True, "stream": True})
+        # noinspection PyProtectedMember
+        socket._sock.setblocking(True)
+        for line in sys.stdin:
+            line = input_line(line)
+            # noinspection PyProtectedMember
+            socket._sock.sendall(line.encode("utf-8"))
+    except (KeyboardInterrupt, EOFError):
+        pass
+    finally:
+        container.stop()
+
+
+def get_sql(json_line: dict) -> Optional[str]:
+    if json_line["method"] == "tools/call":
+        for tool in config["mcp-tools"]:
+            if tool["tool-name"] == json_line["params"]["name"]:
+                return json_line["params"]["arguments"][tool["arg-name"]]
+    return None
+
+
+def input_line(line: str) -> str:
+    json_line = json.loads(line.encode("utf-8"))
+    sql = get_sql(json_line)
+    if sql:
+        result = verify_sql(
+            sql,
+            config["sql-data-guard"],
+            config["sql-data-guard"]["dialect"],
+        )
+        if not result["allowed"]:
+            sys.stderr.write(
+                f"ID: {json_line['id']} Blocked SQL: {sql}\nErrors: {list(result['errors'])}\n"
+            )
+            if result["fixed"]:
+                sys.stderr.write(f"Fixed SQL: {result['fixed']}\n")
+                updated_sql = result["fixed"]
+            else:
+                updated_sql = "SELECT 'Blocked by SQL Data Guard' AS message"
+                if not inject_response:
+                    for error in result["errors"]:
+                        updated_sql += f"\nUNION ALL SELECT '{error}' AS message"
+            if inject_response:
+                result["errors"] = list(result["errors"])
+                if result["fixed"]:
+                    result["fixed"] = result["fixed"].replace("'", "''")
+                errors[json_line["id"]] = result
+            json_line["params"]["arguments"]["query"] = updated_sql
+            line = json.dumps(json_line) + "\n"
+    return line
+
+
+if __name__ == "__main__":
+    config = load_config()
+    inject_response = config["sql-data-guard"]["inject-response"]
+    errors: Dict[int, dict] = {}
+    main()

{sql_data_guard-0.1.7 → sql_data_guard-0.1.9}/src/sql_data_guard/restriction_verification.py

@@ -3,8 +3,8 @@ from typing import List
 import sqlglot
 import sqlglot.expressions as expr
 
-from sql_data_guard.verification_context import VerificationContext
-from sql_data_guard.verification_utils import split_to_expressions
+from .verification_context import VerificationContext
+from .verification_utils import split_to_expressions
 
 
 def verify_restrictions(

{sql_data_guard-0.1.7 → sql_data_guard-0.1.9/src/sql_data_guard.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-data-guard
-Version: 0.1.7
+Version: 0.1.9
 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.7 → sql_data_guard-0.1.9}/src/sql_data_guard.egg-info/SOURCES.txt

@@ -12,6 +12,7 @@ src/sql_data_guard.egg-info/SOURCES.txt
 src/sql_data_guard.egg-info/dependency_links.txt
 src/sql_data_guard.egg-info/requires.txt
 src/sql_data_guard.egg-info/top_level.txt
+src/sql_data_guard/mcpwrapper/mcp_wrapper.py
 src/sql_data_guard/rest/__init__.py
 src/sql_data_guard/rest/sql_data_guard_rest.py
 test/test_duckdb_unit.py

{sql_data_guard-0.1.7 → sql_data_guard-0.1.9}/test/test_sql_guard_validation_unit.py

@@ -1,4 +1,5 @@
 import pytest
+
 from sql_data_guard.restriction_validation import (
     validate_restrictions,
     UnsupportedRestrictionError,