From 2ed1199d48f7207f4a012c04f61e13ac1a8d5154 Mon Sep 17 00:00:00 2001
From: chenjiahe <763432473@qq.com>
Date: 星期四, 16 六月 2022 10:27:44 +0800
Subject: [PATCH] 新增请求安全工具

---
 src/main/java/com/hx/mybatis/aes/springbean/SqlUtils.java |  104 +++++++++++++++++++++++++++++++++++++++++++++-------
 1 files changed, 90 insertions(+), 14 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 90d2228..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,15 +39,88 @@
 
         MySqlStatementParser parser = new MySqlStatementParser(sql);
         SQLSelectStatement sqlStatement = (SQLSelectStatement) parser.parseSelect();
-        //鑾峰彇鏍煎紡鍖栫殑slq璇彞
-        sql = sqlStatement.toString();
 
-        System.out.println("sql:"+sql);
+        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)sqlStatement.getSelect().getQuery() ;
+        SQLSelectQueryBlock sqlSelectQuery = null;
+        boolean b = true;
+        try{
+            sqlSelectQuery = (SQLSelectQueryBlock)sqlStatement.getSelect().getQuery() ;
+        }catch (Exception e){
+            b = false;
+            logger.error("瑙f瀽sql鎶ラ敊锛�"+e.getMessage());
+        }
+        if(!b){
+            return "err";
+        }
+
         StringBuffer out = new StringBuffer() ;
         //鍒涘缓sql瑙f瀽鐨勬爣鍑嗗寲杈撳嚭
         SQLASTOutputVisitor sqlastOutputVisitor = SQLUtils.createFormatOutputVisitor(out , null , JdbcUtils.MYSQL) ;
@@ -119,7 +195,7 @@
         }
 
         //澶勭悊where闇�瑕佸姞瀵嗗緱瀛楁
-        sql = sqlWhere.toString();
+        String sql = sqlWhere.toString();
         if(!StringUtils.isEmpty(sql)){
             Map<String,String> aesKeys = null;
             String aeskey = null;
@@ -140,7 +216,7 @@
                 }else{
                     tableAl = column.getName();
                 }
-                sql = sql.replaceAll("\\b"+tableAl+"\\b","AES_DECRYPT(UNHEX("+tableAl+"),'"+aeskey+"')");
+                sql = sql.replaceAll("((?<!\\.)\\b"+tableAl+"\\b(?!\\.))","AES_DECRYPT(UNHEX("+tableAl+"),'"+aeskey+"')");
             }
         }
         return sqlSelect.toString()+sql;
@@ -249,7 +325,7 @@
                 }else{
                     tableAl = column.getName();
                 }
-                sql = sql.replaceAll("\\b"+tableAl+"\\b","AES_DECRYPT(UNHEX("+tableAl+"),'"+aeskey+"')");
+                sql = sql.replaceAll("((?<!\\.)\\b"+tableAl+"\\b(?!\\.))","AES_DECRYPT(UNHEX("+tableAl+"),'"+aeskey+"')");
             }
         }
         return sqlSelect.toString()+sql;
@@ -413,11 +489,11 @@
                 }else{
                     tableAl = column.getName();
                 }
-                sqlWhere = sqlWhere.replaceAll("\\b"+tableAl+"\\b","AES_DECRYPT(UNHEX("+tableAl+"),'"+aeskey+"')");
+                sqlWhere = sqlWhere.replaceAll("((?<!\\.)\\b"+tableAl+"\\b(?!\\.))","AES_DECRYPT(UNHEX("+tableAl+"),'"+aeskey+"')");
             }
 
         }
-        splicingSql.append(sqlWhere.toString());
+        splicingSql.append(sqlWhere);
         return splicingSql.toString();
     }
 
@@ -491,11 +567,11 @@
                 }else{
                     tableAl = column.getName();
                 }
-                sqlWhere = sqlWhere.replaceAll("\\b"+tableAl+"\\b","AES_DECRYPT(UNHEX("+tableAl+"),'"+aeskey+"')");
+                sqlWhere = sqlWhere.replaceAll("((?<!\\.)\\b"+tableAl+"\\b(?!\\.))","AES_DECRYPT(UNHEX("+tableAl+"),'"+aeskey+"')");
             }
 
         }
-        splicingSql.append(sqlWhere.toString());
+        splicingSql.append(sqlWhere);
         return splicingSql.toString();
     }
 

--
Gitblit v1.8.0