全局数据拦截器兼容分页含sql权限控制

package com.theiavis.workhour.common.Interceptor;


import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.theiavis.security.jwt.UserOperator;
import com.theiavis.workhour.common.utils.StringUtils;
import com.theiavis.workhour.project.dataAuth.bean.DataAuth;
import com.theiavis.workhour.project.dataAuth.mapper.DataAuthMapper;
import com.theiavis.workhour.project.dingtalk.bean.Dept;
import com.theiavis.workhour.project.dingtalk.mapper.DeptMapper;
import com.theiavis.workhour.project.system.role.bean.UserRole;
import com.theiavis.workhour.project.system.role.mapper.UserRoleMapper;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.parser.CCJSqlParserManager;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.ParameterMapping;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.reflection.DefaultReflectorFactory;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Component;

import java.io.StringReader;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
import java.sql.Connection;
import java.util.ArrayList;
import java.util.List;
import java.util.Properties;

/**全局数据拦截器
 * @author lyh
 */
@Component
@Intercepts({

        @Signature(
                type = StatementHandler.class, method = "prepare", args = {
                Connection.class, Integer.class
        })
})
public class DataAuthInterceptor implements Interceptor {
    private static final Logger logger = LoggerFactory.getLogger(DataAuthInterceptor.class);
    /**
     * 表名占位符
     */
    private static final String TABLE_NAME = "#{TABLE_NAME}";
    /**
     * 部门id占位符
     */
    private static final String Permit_Code = "#{PermitCode}";

    @Autowired
    DataAuthMapper dataAuthMapper;


    @Autowired
    @Lazy
    UserOperator userOperator;

    @Autowired
    @Lazy
    UserRoleMapper userRoleMapper;

    @Autowired
    @Lazy
    DeptMapper deptMapper;

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        StatementHandler statementHandler = (StatementHandler) invocation.getTarget();
        //获取拦截下的mapper
        MetaObject metaObject = MetaObject.forObject(statementHandler, SystemMetaObject.DEFAULT_OBJECT_FACTORY,
                SystemMetaObject.DEFAULT_OBJECT_WRAPPER_FACTORY, new DefaultReflectorFactory());
        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");
        MybatisPlusInterceptor dataAuth = getDataAuth(mappedStatement);

        //没有注解直接放行,可根据情况补充更多的业务逻辑
        if (dataAuth == null) {
            return invocation.proceed();
        }
        BoundSql boundSql = statementHandler.getBoundSql();
        String orgSql = boundSql.getSql(); //获取到当前需要被执行的SQL
        //根据指定权限点对原有sql进行修改
        List<ParameterMapping> parameterMappings = boundSql.getParameterMappings();
        Object parameterObject = boundSql.getParameterObject();
        String sql = modifyOrgSql(orgSql, dataAuth, parameterMappings, parameterObject);
        try {
            //通过反射修改sql语句
            Field field = boundSql.getClass().getDeclaredField("sql");
            field.setAccessible(true);
            field.set(boundSql, sql);
            logger.info("修改后的sql" + boundSql.getSql());
        } catch (Exception e) {
            e.printStackTrace();
        }
        return invocation.proceed();
    }


    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
        // TODO Auto-generated method stub
    }

    /**
     * 通过反射获取mapper方法是否加了数据拦截注解
     */
    private MybatisPlusInterceptor getDataAuth(MappedStatement mappedStatement) throws ClassNotFoundException {
        MybatisPlusInterceptor dataAuth = null;
        String id = mappedStatement.getId();
        String className = id.substring(0, id.lastIndexOf("."));
        String methodName = id.substring(id.lastIndexOf(".") + 1);
        final Class<?> cls = Class.forName(className);
        final Method[] methods = cls.getMethods();
        for (Method method : methods) {
            if (method.getName().equals(methodName) && method.isAnnotationPresent(MybatisPlusInterceptor.class)) {
                dataAuth = method.getAnnotation(MybatisPlusInterceptor.class);
                break;
            }
        }
        return dataAuth;
    }

    /**
     * 根据权限点拼装对应sql
     *
     * @return 拼装后的sql
     */
    private String modifyOrgSql(String orgSQql, MybatisPlusInterceptor mybatisPlusInterceptor, List<ParameterMapping> parameterMappings, Object parameterObject) throws JSQLParserException {
        String authSql = dataAuthMapper.selectOne(new QueryWrapper<DataAuth>().eq("auth_role", mybatisPlusInterceptor.roleId())).getAuthSql();
        if (authSql == null || authSql.isEmpty()) {
            return orgSQql;
        }
        //使用CCJSqlParserUtil对sql进行解析
        CCJSqlParserManager parserManager = new CCJSqlParserManager();
        Select select = (Select) parserManager.parse(new StringReader(orgSQql));
        PlainSelect plain = (PlainSelect) select.getSelectBody();
        Table fromItem = (Table) plain.getFromItem();
        //有别名用别名，无别名用表名，防止字段冲突报错
        String aliasTableName = fromItem.getAlias() == null ? fromItem.getName() : fromItem.getAlias().getName();
        //获取个人数据范围部门id
        List<UserRole>deptId=userRoleMapper.selectList(new QueryWrapper<UserRole>().eq("ding_talk_user_id",userOperator.getUser().getUserId()));
        //根据部门id->个人部门id权限码列表
        List<String> arrayList=new ArrayList<String>();
        deptId.forEach(i->{
            Dept deptPermitCode=deptMapper.selectOne(new QueryWrapper<Dept>().eq("id",i.getDeptId()));
            arrayList.add(deptPermitCode.getPermitCode());
        });
        logger.info("获取到用户标识为" + userOperator.getUser().getUserId()+"用户权限为"+userOperator.getUser().getPermitCode());
        String modifSql=new String();
        for (int i=0; i<arrayList.size();i++){
            if (StringUtils.isEmpty(modifSql)){
                modifSql = authSql.replace(TABLE_NAME, aliasTableName).replace(Permit_Code, arrayList.get(i));
                parseSql( modifSql, plain);
            }else {
                modifSql = authSql.replace(TABLE_NAME, aliasTableName).replace(Permit_Code, arrayList.get(i));
                parseSql( modifSql, plain);
            }

        }
        logger.info("修改后sql"+select.toString());
        return select.toString();
    }


    /**
     * 解析sql
     * @param sql
     * @param plain
     * @return
     * @throws JSQLParserException
     */
    public PlainSelect  parseSql(String sql,PlainSelect plain) throws JSQLParserException {
        if (plain.getWhere() == null) {
            plain.setWhere(CCJSqlParserUtil.parseCondExpression(sql));
        } else {

            plain.setWhere(new AndExpression(plain.getWhere(), CCJSqlParserUtil.parseCondExpression(sql)));
        }
        return plain;
    }
}

DROP TABLE IF EXISTS `data_auth`;
CREATE TABLE `data_auth`  (
  `id` int NOT NULL AUTO_INCREMENT,
  `auth_name` varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NULL DEFAULT NULL COMMENT '数据权限名称',
  `auth_sql` varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NULL DEFAULT NULL COMMENT '执行aop权限',
  `auth_role` varchar(255) CHARACTER SET utf8mb4 COLLATE utf8mb4_0900_ai_ci NULL DEFAULT NULL COMMENT '角色权限sys_role->role_id',
  PRIMARY KEY (`id`) USING BTREE
) ENGINE = InnoDB AUTO_INCREMENT = 7 CHARACTER SET = utf8mb4 COLLATE = utf8mb4_0900_ai_ci ROW_FORMAT = DYNAMIC;

SET FOREIGN_KEY_CHECKS = 1;