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

package com.zhgd.mybatis;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.convert.Convert;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSONObject;
import com.baomidou.mybatisplus.extension.plugins.handler.DataPermissionHandler;
import com.zhgd.annotation.DataScope;
import com.zhgd.xmgl.modules.basicdata.enums.SystemUserAccountTypeEnum;
import com.zhgd.xmgl.modules.ocr.service.IOcrBuildLogService;
import com.zhgd.xmgl.modules.video.service.IAiAnalyseHardWareRecordService;
import com.zhgd.xmgl.modules.worker.entity.EnterpriseInfo;
import com.zhgd.xmgl.modules.worker.mapper.EnterpriseInfoMapper;
import com.zhgd.xmgl.modules.worker.service.impl.UserDevAuthorityServiceImpl;
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 com.zhgd.xmgl.util.EnvironmentUtil;
import lombok.extern.slf4j.Slf4j;
import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.*;
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.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Lazy;

import java.util.*;
import java.util.stream.Collectors;

@Slf4j
public class DataScopeHandler implements DataPermissionHandler {
    @Lazy
    @Autowired
    EnvironmentUtil environmentUtil;
    @Lazy
    @Autowired
    IAiAnalyseHardWareRecordService aiAnalyseHardWareRecordService;
    @Lazy
    @Autowired
    private XzSupplierQualificationApplyServiceImpl xzSupplierQualificationApplyService;
    @Lazy
    @Autowired
    private EnterpriseInfoMapper enterpriseInfoMapper;
    @Lazy
    @Autowired
    private UserDevAuthorityServiceImpl userDevAuthorityService;
    @Lazy
    @Autowired
    private UserEnterpriseServiceImpl userEnterpriseService;
    @Lazy
    @Autowired
    private IOcrBuildLogService ocrBuildLogService;

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

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

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

    /**
     * 获取设置的过滤的表名和字段名
     *
     * @return
     */
    private HashMap<String, String> getFieldEnterpriseTables() {
        HashMap<String, String> tables = new HashMap<>(16);
        tables.put("worker_info", "enterprise_id");
        tables.put("team_info", "enterprise_id");
        tables.put("department_info", "enterprise_id");
        tables.put("xz_task_progress_content", "enterprise_id");
        tables.put("xz_task_progress_alarm", "enterprise_id");
        tables.put("xz_task_progress", "enterprise_id");
        tables.put("quality_inspection_record", "enterprise_id");
        tables.put("worker_info_audit_record", "enterprise_id");
        tables.put("xz_material", "enterprise_id");
        tables.put("inspect_task_record", "enterprise_id");
        tables.put("xz_security_inspect_task_record", "enterprise_id");
        tables.put("xz_security_quality_inspection_record", "enterprise_id");
        tables.put("exam_notice", "enterprise_id");
        tables.put("exam_train_record", "enterprise_id");
        tables.put("exam_course_record", "enterprise_id");
        tables.put("xz_worker_safe_watch_alarm", "enterprise_id");
        tables.put("dangerous_engineering_record", "responsibility_company_id");
        tables.put("project_fine_record", "enterprise_id");
        if (!environmentUtil.isBaise()) {
            tables.put("car_info", "enterprise_id");
            tables.put("enterprise_info", "id");
        }
        return tables;
    }

    private HashMap<String, String> getFieldVideoTables() {
        HashMap<String, String> tables = new HashMap<>(16);
        tables.put("video_item", "item_id");
        return tables;
    }

    private HashMap<String, String> getFieldAiTables() {
        HashMap<String, String> tables = new HashMap<>(16);
        if (!environmentUtil.isBaise()) {
            tables.put("ai_analyse_hard_ware_alarm_record", "hardware_id");
        }
        return tables;
    }

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

    private PlainSelect dataScopeFilterByProject(PlainSelect plainSelect, UserInfo user, Object obj) {
        JSONObject jo = (JSONObject) obj;
        Object parameter = jo.get("parameter");
        DataScope ds = jo.getObject("ds", DataScope.class);
        init(plainSelect);
        //expressions
        List<Expression> expressions = new ArrayList<>();
        if (!DataScopeInterceptor.findIgnoreDataScope(parameter, ds)) {
            if (Objects.equals(user.getAccountType(), SystemUserAccountTypeEnum.PROJECT_SUB_ACCOUNT.getValue())) {
                List<String> filterEnterprises = getNeedFilterLeftExpression(plainSelect, getFieldEnterpriseTables(), ds);
                if (CollUtil.isNotEmpty(filterEnterprises)) {
                    List<String> enterpriseIds = userEnterpriseService.getEnterpriseIdsIfSubProject();
                    for (String filterEnterprise : filterEnterprises) {
                        inExpression(filterEnterprise, enterpriseIds, plainSelect);
                    }
                }

                List<String> filterItems = getNeedFilterLeftExpression(plainSelect, getFieldVideoTables(), ds);
                if (CollUtil.isNotEmpty(filterItems)) {
                    List<String> videoItems = userDevAuthorityService.getVideoItemsIfSubProject();
                    for (String item : filterItems) {
                        inExpression(item, videoItems, plainSelect);
                    }
                }

                List<String> filterAis = getNeedFilterLeftExpression(plainSelect, getFieldAiTables(), ds, true);
                if (CollUtil.isNotEmpty(filterAis)) {
                    List<String> videoItems = aiAnalyseHardWareRecordService.getAiAnalyseHardIdsByUserId();
                    if (CollUtil.isEmpty(videoItems)) {
                        videoItems.add("0");
                    }
                    videoItems = videoItems.stream().map(s -> "'" + s + "'").collect(Collectors.toList());
                    for (String filterAi : filterAis) {
                        String sql = StrUtil.format(" ({}.hardware_id in ({}) OR  ({}.quality_region_id in (select distinct quality_region_id from quality_region_to_user where user_id = {}))) ",
                                filterAi, StrUtil.join(",", videoItems), filterAi, SecurityUtils.getUser().getUserId());
                        Expression expression = null;
                        try {
                            expression = CCJSqlParserUtil.parseCondExpression(sql);
                            expressions.add(expression);
                        } catch (JSQLParserException e) {
                            log.error(e.getMessage(), e);
                        }
                    }
                }

                //ocr施工日志施工单位
                HashMap<String, String> ocrBuildLogTables = new HashMap<>();
                ocrBuildLogTables.put("ocr_build_log", "id");
                List<String> filterOcrBuildLogTables = getNeedFilterLeftExpression(plainSelect, ocrBuildLogTables, ds);
                if (CollUtil.isNotEmpty(filterOcrBuildLogTables)) {
                    List<String> ids = ocrBuildLogService.getIdsIfSubProject().stream().map(Convert::toStr).collect(Collectors.toList());
                    ids.add("0");
                    for (String filterEnterprise : filterOcrBuildLogTables) {
                        //( (xxx.id in ()) OR ( xxx.uploader_id = 123))
                        InExpression inExpr = new InExpression();
                        ItemsList itemsList = new ExpressionList(ids.stream().map(StringValue::new).collect(Collectors.toList()));
                        inExpr.setLeftExpression(new Column(filterEnterprise));
                        inExpr.setRightItemsList(itemsList);
                        inExpression(filterEnterprise, ids, plainSelect);
                        // 假设这是你的inExpression方法返回的表达式

                        EqualsTo equalsExpr = new EqualsTo();
                        equalsExpr.setLeftExpression(new Column(StrUtil.subBefore(filterEnterprise, ".", false) + "." + "uploader_id"));
                        equalsExpr.setRightExpression(new LongValue(SecurityUtils.getUser().getUserId()));
                        // 创建OR表达式组合两者
                        OrExpression orExpr = new OrExpression(inExpr, equalsExpr);
                        // 将整个OR表达式添加到WHERE子句中
                        if (plainSelect.getWhere() == null) {
                            plainSelect.setWhere(orExpr);
                        } else {
                            // 如果已有WHERE条件，可能需要用AND连接
                            AndExpression andExpr = new AndExpression(plainSelect.getWhere(), orExpr);
                            plainSelect.setWhere(andExpr);
                        }
                    }
                }

            } else if (Objects.equals(user.getAccountType(), SystemUserAccountTypeEnum.SUPPLIER.getValue())) {
                List<String> filterEnterprises = getNeedFilterLeftExpression(plainSelect, getFieldEnterpriseTables(), ds);
                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);
                    }
                }

                //解析ai预警
                List<String> filterAis = getNeedFilterLeftExpression(plainSelect, getFieldAiTables(), ds, true);
                if (CollUtil.isNotEmpty(filterAis)) {
                    List<String> videoItems = aiAnalyseHardWareRecordService.getAiAnalyseHardIdsByEnterpriseId(id);
                    if (CollUtil.isEmpty(videoItems)) {
                        videoItems.add("0");
                    }
                    videoItems = videoItems.stream().map(s -> "'" + s + "'").collect(Collectors.toList());
                    for (String filterAi : filterAis) {
                        String sql = StrUtil.format(" ({}.hardware_id in ({}) OR  ({}.quality_region_id in (select distinct quality_region_id from quality_region_to_user where user_id = {}))) ",
                                filterAi, StrUtil.join(",", videoItems), filterAi, SecurityUtils.getUser().getUserId());
                        Expression expression = null;
                        try {
                            expression = CCJSqlParserUtil.parseCondExpression(sql);
                            expressions.add(expression);
                        } catch (JSQLParserException e) {
                            log.error(e.getMessage(), e);
                        }
                    }
                }

            }
            if (expressions.size() > 0) {
                Expression dataExpression;
                if (expressions.size() > 1) {
                    //数据权限大于1个，之间用或
                    OrExpression orExpression = new OrExpression(expressions.get(0), expressions.get(1));
                    int i1 = 2;
                    for (int i = i1; 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);
            }

        }


        if (!DataScopeInterceptor.isNotSqlTest()) {
            String sql = " ('1qqq')";
            try {
                Expression expression = CCJSqlParserUtil.parseCondExpression(sql);
                Expression where = plainSelect.getWhere();
                if (where != null) {
                    where = new AndExpression(where, expression);
                } else {
                    where = expression;
                }
                plainSelect.setWhere(where);
            } catch (JSQLParserException e) {
                throw new RuntimeException(e);
            }
        }

        //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;
    }

    /**
     * 获取需要过滤的表别名或加字段
     *
     * @param plainSelect
     * @param tables
     * @param ds
     * @param onlyAlas    false：表别名.字段 true：表别名
     * @return
     */
    private List<String> getNeedFilterLeftExpression(PlainSelect plainSelect, Map<String, String> tables, DataScope ds, boolean onlyAlas) {
        ArrayList<String> rtList = new ArrayList<>();
        String[] dsArr = ds.includeTable();
        if (dsArr.length > 0) {
            Map<String, String> nt = new HashMap<>(16);
            for (String ds1 : dsArr) {
                if (tables.containsKey(ds1)) {
                    nt.put(ds1, tables.get(ds1));
                }
            }
            tables = nt;
        }
        FromItem fromItem = plainSelect.getFromItem();
        if (fromItem instanceof Table) {
            Table table = (Table) fromItem;
            String name = table.getName();
            if (tables.get(name) != null) {
                String aliasName = null;
                Alias alias = table.getAlias();
                if (alias != null) {
                    aliasName = alias.getName();
                }
                String e;
                if (onlyAlas) {
                    e = (StringUtils.isEmpty(aliasName) ? name : aliasName);
                } else {
                    e = (StringUtils.isEmpty(aliasName) ? name : aliasName) + "." + tables.get(name);
                }
                rtList.add(e);
            }
        }
        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.get(joinTable.getName()) != null) {
                        String aliasName = null;
                        Alias alias = joinTable.getAlias();
                        if (alias != null) {
                            aliasName = alias.getName();
                        }
                        String e;
                        if (onlyAlas) {
                            e = (StringUtils.isEmpty(aliasName) ? joinTable.getName() : aliasName);
                        } else {
                            e = (StringUtils.isEmpty(aliasName) ? joinTable.getName() : aliasName) + "." + tables.get(joinTable.getName());
                        }
                        rtList.add(e);
                    }
                }
            }
        }
        return rtList;

    }

    private List<String> getNeedFilterLeftExpression(PlainSelect plainSelect, Map<String, String> tables, DataScope ds) {
        return getNeedFilterLeftExpression(plainSelect, tables, ds, false);
    }

    //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) {
            log.error(e.getMessage(), 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;
    //}
}