[BUG]Fix security bug (#50)
* 0.5.1 Version * fix 0.5.1 schema async bug * fix security bug
This commit is contained in:
@@ -557,23 +557,68 @@ class DorisQueryExecutor:
|
|||||||
"data": None
|
"data": None
|
||||||
}
|
}
|
||||||
|
|
||||||
|
# Import required security modules
|
||||||
|
from .security import DorisSecurityManager, AuthContext, SecurityLevel
|
||||||
|
|
||||||
|
# Create proper auth context with read-only permissions
|
||||||
|
auth_context = AuthContext(
|
||||||
|
user_id=user_id,
|
||||||
|
roles=["read_only_user"], # Restrictive role for MCP interface
|
||||||
|
permissions=["read_data"], # Only read permissions
|
||||||
|
session_id=session_id,
|
||||||
|
security_level=SecurityLevel.INTERNAL
|
||||||
|
)
|
||||||
|
|
||||||
|
# Perform SQL security validation if enabled
|
||||||
|
if hasattr(self.connection_manager, 'config') and hasattr(self.connection_manager.config, 'security'):
|
||||||
|
if self.connection_manager.config.security.enable_security_check:
|
||||||
|
try:
|
||||||
|
security_manager = DorisSecurityManager(self.connection_manager.config)
|
||||||
|
validation_result = await security_manager.validate_sql_security(sql, auth_context)
|
||||||
|
|
||||||
|
if not validation_result.is_valid:
|
||||||
|
self.logger.warning(f"SQL security validation failed for query: {sql[:100]}...")
|
||||||
|
return {
|
||||||
|
"success": False,
|
||||||
|
"error": f"SQL security validation failed: {validation_result.error_message}",
|
||||||
|
"error_type": "security_violation",
|
||||||
|
"blocked_operations": validation_result.blocked_operations,
|
||||||
|
"risk_level": validation_result.risk_level,
|
||||||
|
"data": None,
|
||||||
|
"metadata": {
|
||||||
|
"query": sql,
|
||||||
|
"validation_details": {
|
||||||
|
"blocked_operations": validation_result.blocked_operations,
|
||||||
|
"risk_level": validation_result.risk_level
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
else:
|
||||||
|
self.logger.debug(f"SQL security validation passed for query: {sql[:100]}...")
|
||||||
|
except Exception as security_error:
|
||||||
|
self.logger.error(f"Security validation error: {str(security_error)}")
|
||||||
|
# In case of security validation error, fail safe
|
||||||
|
return {
|
||||||
|
"success": False,
|
||||||
|
"error": f"Security validation system error: {str(security_error)}",
|
||||||
|
"error_type": "security_system_error",
|
||||||
|
"data": None,
|
||||||
|
"metadata": {
|
||||||
|
"query": sql,
|
||||||
|
"security_error": str(security_error)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
else:
|
||||||
|
self.logger.info("SQL security check is disabled in configuration")
|
||||||
|
else:
|
||||||
|
self.logger.warning("Security configuration not found, proceeding without validation")
|
||||||
|
|
||||||
# Add LIMIT if not present and it's a SELECT query
|
# Add LIMIT if not present and it's a SELECT query
|
||||||
if sql.upper().startswith("SELECT") and "LIMIT" not in sql.upper():
|
if sql.upper().startswith("SELECT") and "LIMIT" not in sql.upper():
|
||||||
if sql.endswith(";"):
|
if sql.endswith(";"):
|
||||||
sql = sql[:-1]
|
sql = sql[:-1]
|
||||||
sql = f"{sql} LIMIT {limit}"
|
sql = f"{sql} LIMIT {limit}"
|
||||||
|
|
||||||
# Create auth context for MCP calls
|
|
||||||
class MockAuthContext:
|
|
||||||
def __init__(self):
|
|
||||||
self.user_id = user_id
|
|
||||||
self.roles = ["data_analyst"]
|
|
||||||
self.permissions = ["read_data", "execute_query"]
|
|
||||||
self.session_id = session_id
|
|
||||||
self.security_level = "internal"
|
|
||||||
|
|
||||||
auth_context = MockAuthContext()
|
|
||||||
|
|
||||||
# Create query request
|
# Create query request
|
||||||
query_request = QueryRequest(
|
query_request = QueryRequest(
|
||||||
sql=sql,
|
sql=sql,
|
||||||
@@ -827,9 +872,13 @@ class QueryPerformanceMonitor:
|
|||||||
|
|
||||||
# Unified convenience function for MCP integration
|
# Unified convenience function for MCP integration
|
||||||
async def execute_sql_query(sql: str, connection_manager: DorisConnectionManager, **kwargs) -> Dict[str, Any]:
|
async def execute_sql_query(sql: str, connection_manager: DorisConnectionManager, **kwargs) -> Dict[str, Any]:
|
||||||
"""Execute SQL query - unified convenience function for MCP tools"""
|
"""Execute SQL query - unified convenience function for MCP tools
|
||||||
|
|
||||||
|
This function now includes security validation to ensure safe query execution.
|
||||||
|
All queries are validated against the configured security policies before execution.
|
||||||
|
"""
|
||||||
try:
|
try:
|
||||||
# Create query executor
|
# Create query executor with the connection manager's configuration
|
||||||
executor = DorisQueryExecutor(connection_manager)
|
executor = DorisQueryExecutor(connection_manager)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -839,6 +888,7 @@ async def execute_sql_query(sql: str, connection_manager: DorisConnectionManager
|
|||||||
session_id = kwargs.get("session_id", "mcp_session")
|
session_id = kwargs.get("session_id", "mcp_session")
|
||||||
user_id = kwargs.get("user_id", "mcp_user")
|
user_id = kwargs.get("user_id", "mcp_user")
|
||||||
|
|
||||||
|
# The execute_sql_for_mcp method now includes security validation
|
||||||
result = await executor.execute_sql_for_mcp(
|
result = await executor.execute_sql_for_mcp(
|
||||||
sql=sql,
|
sql=sql,
|
||||||
limit=limit,
|
limit=limit,
|
||||||
@@ -854,5 +904,10 @@ async def execute_sql_query(sql: str, connection_manager: DorisConnectionManager
|
|||||||
return {
|
return {
|
||||||
"success": False,
|
"success": False,
|
||||||
"error": f"Query execution failed: {str(e)}",
|
"error": f"Query execution failed: {str(e)}",
|
||||||
"data": None
|
"error_type": "execution_error",
|
||||||
|
"data": None,
|
||||||
|
"metadata": {
|
||||||
|
"query": sql,
|
||||||
|
"execution_error": str(e)
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user