/*
 * Copyright 2017-2018 Baidu Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.fuxi.javaagent.plugin.checker.local;


import com.baidu.openrasp.TokenGenerator;
import com.fuxi.javaagent.HookHandler;
import com.fuxi.javaagent.config.Config;
import com.fuxi.javaagent.plugin.antlrlistener.TokenizeErrorListener;
import com.fuxi.javaagent.plugin.checker.CheckParameter;
import com.fuxi.javaagent.plugin.checker.js.JsChecker;
import com.fuxi.javaagent.plugin.info.AttackInfo;
import com.fuxi.javaagent.plugin.info.EventInfo;
import com.google.gson.JsonObject;
import org.apache.commons.lang3.StringUtils;

import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

/**
 * Created by tyy on 17-12-20.
 *
 * 检测 sql 语句的 java 版本
 */
public class SqlStatementChecker extends ConfigurableChecker {

    private static final String CONFIG_KEY_SQLI_USER_INPUT = "sqli_userinput";
    private static final String CONFIG_KEY_SQLI_POLICY = "sqli_policy";
    private static final String CONFIG_KEY_STACKED_QUERY = "stacked_query";
    private static final String CONFIG_KEY_NO_HEX = "no_hex";
    private static final String CONFIG_KEY_CONSTANT_COMPARE = "constant_compare";
    private static final String CONFIG_KEY_VERSION_COMMENT = "version_comment";
    private static final String CONFIG_KEY_FUNCTION_BLACKLIST = "function_blacklist";
    private static final String CONFIG_KEY_UNION_NULL = "union_null";

    private static TokenizeErrorListener tokenizeErrorListener = new TokenizeErrorListener();

    @Override
    public List<EventInfo> checkParam(CheckParameter checkParameter) {
        List<EventInfo> result = new LinkedList<EventInfo>();
        String query = (String) checkParameter.getParam("query");

        String message = null;
        String[] tokens = TokenGenerator.tokenize(query, tokenizeErrorListener);
        Map<String, String[]> parameterMap = HookHandler.requestCache.get().getParameterMap();
        JsonObject config = Config.getConfig().getAlgorithmConfig();
        // 算法1: 匹配用户输入
        // 1. 简单识别逻辑是否发生改变
        // 2. 识别数据库管理器
        String action = getActionElement(config, CONFIG_KEY_SQLI_USER_INPUT);
        if (!EventInfo.CHECK_ACTION_IGNORE.equals(action) && action != null) {
            for (Map.Entry<String, String[]> entry : parameterMap.entrySet()) {
                String[] v = entry.getValue();
                String value = v[0];
                if (value.length() <= 15) {
                    continue;
                }
                if (value.length() == query.length() && value.equals(query)) {
                    message = "算法2: WebShell - 数据库管理器";
                    break;
                }
                if (!query.contains(value)) {
                    continue;
                }
                String[] tokens2 = TokenGenerator.tokenize(query.replace(value, ""), tokenizeErrorListener);
                if (tokens != null) {
                    if (tokens.length - tokens2.length > 2) {
                        message = "算法1: 数据库查询逻辑发生改变";
                        break;
                    }
                }
            }
        }
        if (message != null) {
            result.add(AttackInfo.createLocalAttackInfo(checkParameter, action,
                    message, 90));
        } else {
            // 算法2: SQL语句策略检查（模拟SQL防火墙功能）
            action = getActionElement(config, CONFIG_KEY_SQLI_POLICY);
            if (!EventInfo.CHECK_ACTION_IGNORE.equals(action)) {
                int i = -1;
                if (tokens != null) {
                    HashSet<String> modules = getJsonArrayAsSet(config, CONFIG_KEY_SQLI_POLICY, "feature");
                    for (String token : tokens) {
                        i++;
                        if (!StringUtils.isEmpty(token)) {
                            String lt = token.toLowerCase();
                            if (lt.equals("select") && modules.contains(CONFIG_KEY_UNION_NULL)) {
                                int nullCount = 0;
                                // 寻找连续的逗号、NULL或者数字
                                for (int j = i + 1; j < tokens.length && j < i + 6; j++) {
                                    if (tokens[j].equals(",") || tokens[j].equals("null") || StringUtils.isNumeric(tokens[j])) {
                                        nullCount++;
                                    } else {
                                        break;
                                    }
                                }

                                // NULL,NULL,NULL == 5个token
                                // 1,2,3          == 5个token
                                if (nullCount >= 5) {
                                    message = "UNION-NULL 方式注入 - 字段类型探测";
                                    break;
                                }
                                continue;
                            }
                            if (lt.equals(";") && i != tokens.length - 1 && modules.contains(CONFIG_KEY_STACKED_QUERY)) {
                                message = "禁止多语句查询";
                                break;
                            } else if (lt.startsWith("0x") && modules.contains(CONFIG_KEY_NO_HEX)) {
                                message = "禁止16进制字符串";
                                break;
                            } else if (lt.startsWith("/*!") && modules.contains(CONFIG_KEY_VERSION_COMMENT)) {
                                message = "禁止MySQL版本号注释";
                                break;
                            } else if (i > 0 && i < tokens.length - 1 && (lt.equals("xor")
                                    || lt.charAt(0) == '<'
                                    || lt.charAt(0) == '>'
                                    || lt.charAt(0) == '=') && modules.contains(CONFIG_KEY_CONSTANT_COMPARE)) {
                                String op1 = tokens[i - 1];
                                String op2 = tokens[i + 1];
                                if (StringUtils.isNumeric(op1) && StringUtils.isNumeric(op2)) {
                                    try {
                                        if (Double.parseDouble(op1) < 10 && Double.parseDouble(op2) < 10) {
                                            continue;
                                        }
                                    } catch (Exception e) {
                                        // ignore
                                    }
                                    message = "禁止常量比较操作";
                                    break;
                                }
                            } else if (i > 0 && tokens[i].indexOf('(') == 0
                                    && modules.contains(CONFIG_KEY_FUNCTION_BLACKLIST)) {
                                // FIXME: 可绕过，暂时不更新
                                HashSet<String> funBlackList = getJsonArrayAsSet(config, CONFIG_KEY_SQLI_POLICY, "function_blacklist");
                                if (funBlackList.contains(tokens[i - 1])) {
                                    message = "禁止执行敏感函数: " + tokens[i - 1];
                                    break;
                                }
                            }
                        }
                    }
                }
                if (message != null) {
                    result.add(AttackInfo.createLocalAttackInfo(checkParameter, action,
                            message, 100));
                }
            }
        }

        // js 插件检测
        List<EventInfo> jsResults = new JsChecker().checkParam(checkParameter);
        if (jsResults != null && jsResults.size() > 0) {
            result.addAll(jsResults);
        }
        return result;
    }

}

@@ -6689,18 +6689,18 @@
 1) %3C 10 
-&&
+%7C%7C
  Double.
@@ -7023,17 +7023,40 @@
 %E7%A6%81%E6%AD%A2%E5%B8%B8%E9%87%8F%E6%AF%94%E8%BE%83%E6%93%8D%E4%BD%9C
-%22
+: %22 + op1 + %22 vs %22 + op2
 ;%0A      

modify the logic of java sql detection