package cc.smtweb.framework.core.mvc.filter;

import javax.servlet.ServletException;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.Enumeration;
import java.util.Map;
import java.util.Set;


public class XssHttpRequestWrapper extends HttpServletRequestWrapper {

    private XssSecurityConfig config;

    /**
     * 封装http请求
     *
     * @param request
     */
    public XssHttpRequestWrapper(HttpServletRequest request, XssSecurityConfig config) {
        super(request);
        this.config = config;
    }

    @Override
    public String getHeader(String name) {
        String value = super.getHeader(name);
        // 若开启特殊字符替换，对特殊字符进行替换
        if (config.isReplace()) {
            config.securityReplace(name);
        }
        return value;
    }

    @Override
    public String getParameter(String name) {
        String value = super.getParameter(name);
        // 若开启特殊字符替换，对特殊字符进行替换
        if (config.isReplace()) {
            config.securityReplace(name);
        }
        return value;
    }

    /**
     * 没有违规的数据，就返回false;
     *
     * @return
     */
    @SuppressWarnings("unchecked")
    private boolean checkHeader() {
        Enumeration<String> headerParams = this.getHeaderNames();
        while (headerParams.hasMoreElements()) {
            String headerName = headerParams.nextElement();
            String headerValue = this.getHeader(headerName);
            if (config.matches(headerValue)) {
                //UtilLogger.debug("XSS校验：检测到Header攻击：[" + headerValue + "]");
                return true;
            }
        }
        return false;
    }

    /**
     * 没有违规的数据，就返回false;
     *
     * @return
     */
    @SuppressWarnings("unchecked")
    private boolean checkParameter() {
        Map<String, String[]> submitParams = this.getParameterMap();
        Set<String> submitNames = submitParams.keySet();
        for (String submitName : submitNames) {
            Object submitValues = submitParams.get(submitName);
            if (submitValues instanceof String) {
                if (config.matches((String) submitValues)) {
                    //UtilLogger.debug("XSS校验：检测到Parameter攻击：[" + submitValues + "]");
                    return true;
                }
            } else if (submitValues instanceof String[]) {
                for (String submitValue : (String[]) submitValues) {
                    if (config.matches((String) submitValue)) {
                        //UtilLogger.debug("XSS校验：检测到Parameter攻击：[" + submitValue + "]");
                        return true;
                    }
                }
            }
        }
        return false;
    }


    /**
     * 没有违规的数据，就返回false;
     * 若存在违规数据，根据配置信息判断是否跳转到错误页面
     *
     * @param response
     * @return
     * @throws IOException
     * @throws ServletException
     */
    public boolean validateParameter(HttpServletResponse response) throws ServletException, IOException {
        // 开始header校验，对header信息进行校验
        if (config.isCheckHeader()) {
            return this.checkHeader();
        }
        // 开始parameter校验，对parameter信息进行校验
        if (config.isCheckParameter()) {
            return this.checkParameter();
        }
        return false;
    }
}
