wisdomisite-java/src/main/java/com/zhgd/mybatis/DataScopeHandler.java

package com.zhgd.mybatis;

import cn.hutool.core.util.StrUtil;
import com.baomidou.mybatisplus.extension.plugins.handler.DataPermissionHandler;
import com.zhgd.xmgl.modules.worker.entity.EnterpriseInfo;
import com.zhgd.xmgl.modules.worker.mapper.EnterpriseInfoMapper;
import com.zhgd.xmgl.modules.worker.service.impl.UserEnterpriseServiceImpl;
import com.zhgd.xmgl.modules.xz.service.impl.XzSupplierQualificationApplyServiceImpl;
import com.zhgd.xmgl.security.entity.UserInfo;
import com.zhgd.xmgl.security.util.SecurityUtils;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.Alias;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.Parenthesis;
import net.sf.jsqlparser.expression.StringValue;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.expression.operators.relational.ItemsList;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.select.FromItem;
import net.sf.jsqlparser.statement.select.Join;
import net.sf.jsqlparser.statement.select.PlainSelect;
import org.apache.commons.collections.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Lazy;

import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

@Slf4j
public class DataScopeHandler implements DataPermissionHandler {
    @Lazy
    @Autowired
    private XzSupplierQualificationApplyServiceImpl xzSupplierQualificationApplyService;
    @Lazy
    @Autowired
    private EnterpriseInfoMapper enterpriseInfoMapper;
    //
    //@Lazy
    //@Autowired
    //private IProjectService projectService;
    //
    //@Lazy
    //@Autowired
    //private IEnterpriseService enterpriseService;
    //
    //@Lazy
    //@Autowired
    //private IEngineeringService engineeringService;
    //
    //@Lazy
    //@Autowired
    //private ISystemUserDataScopeService systemUserDataScopeService;

    @Lazy
    @Autowired
    private UserEnterpriseServiceImpl userEnterpriseService;


    @Override
    public Expression getSqlSegment(Expression where, String mappedStatementId) {
        return null;
    }

    protected String getEngineeringSn() {
        return "engineering_sn";
    }

    protected String getProjectSn() {
        return "project_sn";
    }

    private String getEnterpriseSn() {
        return "enterprise_sn";
    }

    public PlainSelect getSqlSegment(PlainSelect plainSelect, Object obj) {
        UserInfo user = SecurityUtils.getUser();
        if (user == null) {
            return plainSelect;
        }
        Integer type = Integer.parseInt(obj.toString());
        if (type == 1) {
            return dataScopeFilterByProject(plainSelect, user);
        }
        // else if (type == 2 && user.getAccountType() == 4) {
        //    return dataScopeFilterByProject(plainSelect, user);
        //} else if (type == 2) {
        //    return dataScopeFilterByEnt(plainSelect, user);
        //}
        return plainSelect;
    }

    private PlainSelect dataScopeFilterByProject(PlainSelect plainSelect, UserInfo user) {
        init(plainSelect);
        List<Expression> expressions = new ArrayList<>();
        if (user.getAccountType() == 6) {
            List<String> enterpriseIds = userEnterpriseService.getEnterpriseIdsIfSubProject();
            List<String> filterEnterprises = getNeedFilterLeftExpression(plainSelect, getFieldEnterpriseTables(), "enterprise_id");
            for (String filterEnterprise : filterEnterprises) {
                inExpression(filterEnterprise, enterpriseIds, plainSelect);
            }
        } else if (user.getAccountType() == 11) {
            List<String> filterEnterprises = getNeedFilterLeftExpression(plainSelect, getFieldEnterpriseTables(), "enterprise_id");
            EnterpriseInfo ei = enterpriseInfoMapper.getXzSupplierByUserId(SecurityUtils.getUser().getUserId());
            Long id;
            if (ei == null) {
                id = -1L;
            } else {
                id = ei.getId();
            }
            for (String filterEnterprise : filterEnterprises) {
                String sql = StrUtil.format(" ({} = {} OR {} IN ( SELECT DISTINCT t.enterprise_id FROM " +
                                "(SELECT t.id FROM project_enterprise t WHERE t.enterprise_id = {}) t2 join project_enterprise t on find_in_set( t2.id, ancestors ) )) ",
                        filterEnterprise, id, filterEnterprise, id);
                Expression expression = null;
                try {
                    expression = CCJSqlParserUtil.parseCondExpression(sql);
                    expressions.add(expression);
                } catch (JSQLParserException e) {
                    log.error(e.getMessage(), e);
                }
            }
            Expression dataExpression;
            if (expressions.size() > 1) {
                //数据权限大于1个，之间用或
                OrExpression orExpression = new OrExpression(expressions.get(0), expressions.get(1));
                for (int i = 2; i < expressions.size(); i++) {
                    orExpression = new OrExpression(orExpression, expressions.get(i));
                }
                // 数据权限使用单独的括号 防止与其他条件冲突
                dataExpression = new Parenthesis(orExpression);
            } else {
                dataExpression = expressions.get(0);
            }
            Expression where = plainSelect.getWhere();
            if (where != null) {
                where = new AndExpression(where, dataExpression);
            } else {
                where = dataExpression;
            }
            plainSelect.setWhere(where);
        }


        //List<String> scopeIds = systemUserDataScopeService.list(Wrappers.<SystemUserDataScope>lambdaQuery().eq(SystemUserDataScope::getUserId, user.getUserId()))
        //        .stream().map(u -> u.getRelevanceId()).collect(Collectors.toList());
        //if (user.getAccountType() == 4) {
        //    if (scopeIds.size() > 0) {
        //        inExpression(getAliasColumn(plainSelect, getEngineeringSn()), scopeIds, plainSelect);
        //    } else {
        //        equalsTo(getAliasColumn(plainSelect, getProjectSn()), user.getSn(), plainSelect);
        //    }
        //}
        //if (user.getAccountType() == 3) {
        //    List<String> engineeringSns = scopeIds.size() > 0 ? scopeIds : engineeringService.getSnListByEnterprise(user.getSn());
        //    if (engineeringSns.size() == 0) {
        //        engineeringSns.add("0");
        //    }
        //    inExpression(getAliasColumn(plainSelect, getEngineeringSn()), engineeringSns, plainSelect);
        //}
        //if (user.getAccountType() == 2) {
        //    List<String> projectSns = projectService.getSnListForGov(user.getSn());
        //    if (projectSns.size() == 0) {
        //        projectSns.add("0");
        //    }
        //    inExpression(getAliasColumn(plainSelect, getProjectSn()), projectSns, plainSelect);
        //}
        return plainSelect;
    }

    @NotNull
    private List<String> getFieldEnterpriseTables() {
        List<String> tables = new ArrayList<>();
        tables.add("worker_info");
        return tables;
    }

    private List<String> getNeedFilterLeftExpression(PlainSelect plainSelect, List<String> tables, String field) {
        ArrayList<String> rtList = new ArrayList<>();
        FromItem fromItem = plainSelect.getFromItem();
        if (fromItem instanceof Table) {
            Table table = (Table) fromItem;
            String name = table.getName();
            if (tables.contains(name)) {
                String aliasName = null;
                Alias alias = table.getAlias();
                if (alias != null) {
                    aliasName = alias.getName();
                }
                rtList.add((StringUtils.isEmpty(aliasName) ? name : aliasName) + "." + field);
            }
        }
        List<Join> joins = plainSelect.getJoins();
        //主表不是user表
        if (!CollectionUtils.isEmpty(joins)) {
            //判断join的表里有没有user表
            for (Join join : joins) {
                // 判断join里面是否存在user表，不存在则新增
                FromItem rightItem = join.getRightItem();
                if (rightItem instanceof Table) {
                    Table joinTable = (Table) rightItem;
                    if (tables.contains(joinTable.getName())) {
                        String aliasName = null;
                        Alias alias = joinTable.getAlias();
                        if (alias != null) {
                            aliasName = alias.getName();
                        }
                        rtList.add((StringUtils.isEmpty(aliasName) ? joinTable.getName() : aliasName) + "." + field);
                    }
                }
            }
        }
        return rtList;
    }

    //private PlainSelect dataScopeFilterByEnt(PlainSelect plainSelect, SecurityUser user) {
    //    init(plainSelect);
    //    if (user.getAccountType() == 3) {
    //        equalsTo(getAliasColumn(plainSelect, getEnterpriseSn()), user.getSn(), plainSelect);
    //    }
    //    if (user.getAccountType() == 2) {
    //        List<String> entSns = enterpriseService.getSnListForGov(user.getSn());
    //        if (entSns.size() == 0) {
    //            entSns.add("0");
    //        }
    //        inExpression(getAliasColumn(plainSelect, getEnterpriseSn()), entSns, plainSelect);
    //    }
    //    return plainSelect;
    //}

    private void equalsTo(String leftExpression, String rightExpression, PlainSelect plainSelect) {
        EqualsTo equalsTo = new EqualsTo();
        equalsTo.setLeftExpression(new Column(leftExpression));
        equalsTo.setRightExpression(new StringValue(rightExpression));
        AndExpression andExpression = new AndExpression(plainSelect.getWhere(), equalsTo);
        plainSelect.setWhere(andExpression);
    }

    private void inExpression(String leftExpression, List<String> rightExpression, PlainSelect plainSelect) {
        InExpression inExpression = new InExpression();
        ItemsList itemsList = new ExpressionList(rightExpression.stream().map(StringValue::new).collect(Collectors.toList()));
        inExpression.setLeftExpression(new Column(leftExpression));
        inExpression.setRightItemsList(itemsList);
        AndExpression andExpression = new AndExpression(plainSelect.getWhere(), inExpression);
        plainSelect.setWhere(andExpression);
    }

    private String getAliasColumn(PlainSelect plainSelect, String sn) {
        FromItem fromItem = plainSelect.getFromItem();
        Alias alias = fromItem.getAlias();
        StringBuilder prefix = new StringBuilder();
        if (alias != null) {
            prefix.append(alias).append(".");
        }
        return prefix.append(sn).toString();
    }

    private void init(PlainSelect plainSelect) {
        Expression envCondition = null;
        try {
            envCondition = CCJSqlParserUtil.parseCondExpression("1 == 1");
        } catch (JSQLParserException e) {
            throw new RuntimeException(e);
        }
        Expression where = plainSelect.getWhere();
        if (where == null) {
            plainSelect.setWhere(envCondition);
        }
    }

    //protected void addParam(Insert insert, Object obj){
    //    SecurityUser user = SecurityUtil.getUser();
    //    if (user == null) {
    //        return;
    //    }
    //    List<Column> columns = insert.getColumns();
    //    if (CollectionUtils.isEmpty(columns)) {
    //        // 针对不给列名的insert 不处理
    //        return;
    //    }
    //    if (user.getAccountType() == 4 || user.getAccountType() == 3) {
    //        if (!columns.stream().anyMatch(c -> c.toString().equals(getProjectSn()))) {
    //            columns.add(new Column(getProjectSn()));
    //            if (insert.getItemsList() != null) {
    //                ItemsList itemsList = insert.getItemsList();
    //                if (itemsList instanceof MultiExpressionList) {
    //                    ((MultiExpressionList) itemsList).getExprList().forEach(el -> el.getExpressions().add(getProjectSn(obj)));
    //                } else {
    //                    ((ExpressionList) itemsList).getExpressions().add(getProjectSn(obj));
    //                }
    //            } else {
    //                throw ExceptionUtils.mpe("Failed to process multiple-table update, please exclude the tableName or statementId");
    //            }
    //        }
    //    }
    //}

    //protected Expression getProjectSn(Object engineeringSn){
    //    SecurityUser user = SecurityUtil.getUser();
    //    if (user.getAccountType() == 4) {
    //        return new StringValue(user.getSn());
    //    }
    //    if (user.getAccountType() == 3) {
    //        return new StringValue(engineeringService.getOne(Wrappers.<Engineering>lambdaQuery().eq(Engineering::getEngineeringSn, engineeringSn)).getProjectSn());
    //    }
    //    return null;
    //}
}