From 42c9e166a8835ad1986d4fb55299acb64623157c Mon Sep 17 00:00:00 2001
From: chenjiahe <763432473@qq.com>
Date: 星期四, 05 五月 2022 13:36:09 +0800
Subject: [PATCH] 修复MYSQL查询语句不兼容问题,特殊union

---
 src/main/java/com/hx/mybatis/aes/springbean/SqlUtils.java |   88 ++++++++++++++++++++++++++++++++++++++++----
 1 files changed, 80 insertions(+), 8 deletions(-)

diff --git a/src/main/java/com/hx/mybatis/aes/springbean/SqlUtils.java b/src/main/java/com/hx/mybatis/aes/springbean/SqlUtils.java
index 3da69dc..8d978db 100644
--- a/src/main/java/com/hx/mybatis/aes/springbean/SqlUtils.java
+++ b/src/main/java/com/hx/mybatis/aes/springbean/SqlUtils.java
@@ -2,7 +2,6 @@
 
 import com.alibaba.druid.sql.SQLUtils;
 import com.alibaba.druid.sql.ast.SQLExpr;
-import com.alibaba.druid.sql.ast.SQLObject;
 import com.alibaba.druid.sql.ast.SQLStatement;
 import com.alibaba.druid.sql.ast.statement.*;
 import com.alibaba.druid.sql.dialect.mysql.ast.statement.MySqlDeleteStatement;
@@ -15,8 +14,9 @@
 import com.alibaba.druid.util.JdbcConstants;
 import com.alibaba.druid.util.JdbcUtils;
 import com.hx.util.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
-import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -26,8 +26,11 @@
  * @author CJH 2022-01-12
  */
 public class SqlUtils {
+    //log4j鏃ュ織
+    private static Logger logger = LoggerFactory.getLogger(SqlUtils.class.getName());
 
-    /**鏌ヨ鍔犲瘑鏁版嵁澶勭悊锛屽彧瀵规煡璇㈠仛澶勭悊锛宻elect杩斿洖涓嶅仛澶勭悊
+
+    /**鏌ヨ鍔犲瘑鏁版嵁澶勭悊锛屽彧瀵规煡璇㈠仛澶勭悊
      * @param sql sql璇彞
      * @param aesKeysTable aes绉橀挜
      * @return
@@ -36,17 +39,86 @@
 
         MySqlStatementParser parser = new MySqlStatementParser(sql);
         SQLSelectStatement sqlStatement = (SQLSelectStatement) parser.parseSelect();
-        //鑾峰彇鏍煎紡鍖栫殑slq璇彞
-        sql = sqlStatement.toString();
+
+        SQLSelect sqlSelect = sqlStatement.getSelect();
+        if (sqlSelect.getQuery() instanceof SQLSelectQueryBlock) {
+            // 闈瀠nion鐨勬煡璇㈣鍙�
+            return selectSqlRoutine( sqlStatement,aesKeysTable);
+        } else if (sqlSelect.getQuery() instanceof SQLUnionQuery) {
+            // union鐨勬煡璇㈣鍙�
+            return selectSqlUnion( sql, sqlStatement, aesKeysTable);
+        }else {
+            return selectSqlRoutine( sqlStatement,aesKeysTable);
+        }
+    }
+
+    /**鏌ヨ鍔犲瘑鏁版嵁澶勭悊锛屽彧瀵规煡璇㈠仛澶勭悊锛宻elect杩斿洖涓嶅仛澶勭悊锛圲nion鐗规畩璇彞锛�
+     * @param sql sql璇彞
+     * @param aesKeysTable aes绉橀挜
+     * @return
+     */
+    public static String selectSqlUnion(String sql,SQLSelectStatement sqlStatement,Map<String,Map<String,String>> aesKeysTable){
+
+        //鑾峰彇琛ㄥ拰鍒悕
+        ExportTableAliasVisitor visitorTable = new ExportTableAliasVisitor();
+        sqlStatement.accept(visitorTable);
+        Map<String,String> tableMaps = visitorTable.getTableMap();
+
+        //鑾峰彇鎵�鏈夌殑瀛楁
+        MySqlSchemaStatVisitor visitor = new MySqlSchemaStatVisitor();
+        sqlStatement.accept(visitor);
+        //閬嶅巻鎵�鏈夊瓧娈�
+        Collection<TableStat.Column> columns= visitor.getColumns();
+
+        //澶勭悊闇�瑕佸姞瀵嗗緱瀛楁
+
+        if(!StringUtils.isEmpty(sql)){
+            Map<String,String> aesKeys = null;
+            String aeskey = null;
+            //鎶婂墿涓嬬殑鎷兼帴涓婃潵
+            String tableAl = null;
+            for(TableStat.Column column:columns){
+                aesKeys= aesKeysTable.get(column.getTable());
+                if(aesKeys == null){
+                    continue;
+                }
+                aeskey = aesKeys.getOrDefault(column.getName(),null);
+                if(StringUtils.isEmpty(aeskey)){
+                    continue;
+                }
+                tableAl = tableMaps.get(column.getTable());
+                if(!StringUtils.isEmpty(tableAl)){
+                    tableAl = tableAl+"."+column.getName();
+                }else{
+                    tableAl = column.getName();
+                }
+                sql = sql.replaceAll("((?<!\\.)\\b"+tableAl+"\\b(?!\\.))","AES_DECRYPT(UNHEX("+tableAl+"),'"+aeskey+"')");
+            }
+        }
+        return sql;
+    }
+
+
+    /**鏌ヨ鍔犲瘑鏁版嵁澶勭悊锛屽彧瀵规煡璇㈠仛澶勭悊锛宻elect杩斿洖涓嶅仛澶勭悊锛堝父瑙勮鍙ワ級
+     * @param sqlStatement sql璇彞
+     * @param aesKeysTable aes绉橀挜
+     * @return
+     */
+    public static String selectSqlRoutine(SQLSelectStatement sqlStatement,Map<String,Map<String,String>> aesKeysTable){
 
         //瑙f瀽select鏌ヨ
         //SQLSelect sqlSelect = sqlStatement.getSelect()
         //鑾峰彇sql鏌ヨ鍧�
-        SQLSelectQueryBlock sqlSelectQuery;
+        SQLSelectQueryBlock sqlSelectQuery = null;
+        boolean b = true;
         try{
             sqlSelectQuery = (SQLSelectQueryBlock)sqlStatement.getSelect().getQuery() ;
         }catch (Exception e){
-            return sql;
+            b = false;
+            logger.error("瑙f瀽sql鎶ラ敊锛�"+e.getMessage());
+        }
+        if(!b){
+            return "err";
         }
 
         StringBuffer out = new StringBuffer() ;
@@ -123,7 +195,7 @@
         }
 
         //澶勭悊where闇�瑕佸姞瀵嗗緱瀛楁
-        sql = sqlWhere.toString();
+        String sql = sqlWhere.toString();
         if(!StringUtils.isEmpty(sql)){
             Map<String,String> aesKeys = null;
             String aeskey = null;

--
Gitblit v1.8.0