package com.xxx;

import cn.hutool.core.util.EscapeUtil;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {

    public static final String HEADER_XSS_FILTER = "xssFilter";

    public static final String[] EMPTY_STRING_ARRAY = new String[0];

    /**
     * 默认全角转义
     */
    public static final String XSS_FILTER_FULL_ANGLE = "1";
    /**
     * html转义
     */
    public static final String XSS_FILTER_ESCAPE_HTML = "2";


    /**
     * Constructs a request object wrapping the given request.
     *
     * @param request The request to wrap
     * @throws IllegalArgumentException if the request is null
     */
    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
    }


    /**
     * 全角字符转义
     *
     * @param str 待转义
     * @return 转义结果
     * @author wulingming
     * @date 2021/9/9 17:07
     **/
    public static String escapeFullAngle(String str) {
        if (str == null || str.length() == 0) {
            return str;
        } else {
            return str.replace("<", "《").replace(">", "》");
        }
    }

    public static String escapeHtml4(String str) {
        if (str == null || str.length() == 0) {
            return str;
        } else {
            return EscapeUtil.escapeHtml4(str);
        }
    }

    private static String escape(String value) {
        return escape(value, null);
    }

    private static String escape(String value, String headerXssFilter) {
        if (XSS_FILTER_ESCAPE_HTML.equals(headerXssFilter)) {
            //html转义
            return escapeHtml4(value);
        } else {
            //默认全角转义
            return escapeFullAngle(value);
        }
    }


    @Override
    public String getParameter(String name) {
        String value = super.getParameter(name);
        if (value == null || value.length() == 0) {
            return null;
        }
        String headerXssFilter = super.getHeader(HEADER_XSS_FILTER);
        return escape(value, headerXssFilter);
    }


    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        if (values == null || values.length == 0) {
            return values;
        }

        String headerXssFilter = super.getHeader(HEADER_XSS_FILTER);

        return Stream.of(values).map(value -> escape(value, headerXssFilter)).collect(Collectors.toList())
                .toArray(EMPTY_STRING_ARRAY);
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> parameterMap = super.getParameterMap();
        if (parameterMap == null || parameterMap.isEmpty()) {
            return parameterMap;
        }

        String headerXssFilter = super.getHeader(HEADER_XSS_FILTER);

        return parameterMap.keySet().stream().filter(Objects::nonNull)
                .collect(Collectors.toMap(key -> key, key -> {
                    String[] values = parameterMap.get(key);
                    if (values == null || values.length == 0) {
                        return EMPTY_STRING_ARRAY;
                    } else {
                        return Stream.of(values).map(value -> escape(value, headerXssFilter))
                                .collect(Collectors.toList()).toArray(EMPTY_STRING_ARRAY);
                    }
                }));
    }
}

/**
 * xss攻击防范处理
 * @author wulingming
 * @date 2021/9/9 16:32
 **/
@WebFilter(urlPatterns = "/*", filterName = "xssFilter")
public class XssFilter implements Filter {


    @Override
    public void init(FilterConfig filterConfig) throws ServletException {

    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response,
            FilterChain chain) throws IOException, ServletException {
        XssHttpServletRequestWrapper req=new XssHttpServletRequestWrapper((HttpServletRequest)request);
        chain.doFilter(req,response);

    }
}

@ServletComponentScan