package com.njcn.web.filter;

import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.databind.JsonSerializer;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializerProvider;
import com.njcn.common.utils.LogUtil;
import lombok.extern.slf4j.Slf4j;
import net.sf.json.JSONArray;
import net.sf.json.JSONObject;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.lang3.StringEscapeUtils;
import org.apache.commons.lang3.StringUtils;
import org.owasp.validator.html.*;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;
import java.net.URLDecoder;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * @author hongawen
 * @version 1.0.0
 * @date 2021年05月20日 16:34
 */
@Slf4j
public class XssRequestWrapper extends HttpServletRequestWrapper {

    private final static String[] WHITE_PARAMETER_NAME = {"password", "mxContent", "docContent", "bgImage","fileContent","flowableXml","bpmnXml","fields","definition","thumbnail","image"};


    public XssRequestWrapper(HttpServletRequest request) {
        super(request);
    }

    private static Policy policy = null;

    /**
     * 获取策略文件，直接使用jar中自带的ebay策略文件
     */
    private static final InputStream INPUT_STREAM = XssRequestWrapper.class.getClassLoader()
            .getResourceAsStream("antisamy-ebay.xml");


    static {
        try {
            // 使用静态代码块处理策略对象的创建
            policy = Policy.getInstance(INPUT_STREAM);
        } catch (PolicyException e) {
            e.printStackTrace();
        }
    }


//    static {
//        try {
//            // 获取策略文件路径，策略文件需要放到项目的classpath下
//            String antiSamyPath = Objects
//                    .requireNonNull(XssRequestWrapper.class.getClassLoader().getResource("antisamy-ebay.xml")).getFile();
//            log.info(antiSamyPath);
//            // 获取的文件路径中有空格时，空格会被替换为%20，在new一个File对象时会出现找不到路径的错误
//            // 对路径进行解码以解决该问题
//            antiSamyPath = URLDecoder.decode(antiSamyPath, "utf-8");
//            log.info(antiSamyPath);
//            // 指定策略文件
//            policy = Policy.getInstance(antiSamyPath);
//        } catch (UnsupportedEncodingException | PolicyException e) {
//            e.printStackTrace();
//        }
//    }


    /**
     * 使用AntiSamy进行过滤数据
     */
    private String xssClean(String html) {
        String cleanHtml = "";
        try {
            AntiSamy antiSamy = new AntiSamy();
            CleanResults scan = antiSamy.scan(html, policy);
            cleanHtml = scan.getCleanHTML();
            cleanHtml = cleanHtml.replace("\n", "");
            // 对转义的HTML特殊字符（<、>、"等）进行反转义，因为AntiSamy调用scan方法时会将特殊字符转义
            cleanHtml = StringEscapeUtils.unescapeHtml4(cleanHtml);
        } catch (ScanException | PolicyException e) {
            log.error("使用AntiSamy进行xss过滤数据异常，过滤内容：{}，异常为：{}", html, e.getMessage());
        }
        return cleanHtml;
    }

    /**
     * 过滤请求头
     *
     * @param name 参数名
     * @return 参数值
     */
    @Override
    public String getHeader(String name) {
        String header = super.getHeader(name);
        // 如果Header为空，则直接返回，否则进行清洗
        return StringUtils.isBlank(header) ? header : xssClean(header);
    }

    /**
     * 重写处理请求参数的方法
     *
     * @param name 参数名
     */
    @Override
    public String getParameter(String name) {
        String parameter = super.getParameter(name);
        if (StringUtils.isEmpty(parameter)) {
            return null;
        }
        LogUtil.njcnDebug(log, "使用AntiSamy进行过滤清理，过滤清理之前的数据：{}", parameter);
        parameter = xssClean(parameter);
        LogUtil.njcnDebug(log, "使用AntiSamy进行过滤清理，过滤清理之后的数据：{}", parameter);
        return parameter;
    }


    /**
     * 重写处理请求参数的方法
     *
     * @param name 参数名
     */
    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        // 判断参数有值，如果没有值，直接返回
        if (values == null) {
            return null;
        }
        return Stream.of(values).map(value -> {
            LogUtil.njcnDebug(log, "使用AntiSamy进行过滤清理，过滤清理之前的数据：{}", value);
            value = xssClean(value);
            LogUtil.njcnDebug(log, "使用AntiSamy进行过滤清理，过滤清理之后的数据：{}", value);
            return value;
        }).toArray(String[]::new);
    }

    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> requestMap = super.getParameterMap();
        requestMap.forEach((key, value) -> {
            for (int i = 0; i < value.length; i++) {
                LogUtil.njcnDebug(log, "使用AntiSamy进行过滤清理，过滤清理之前的数据：{}", value[i]);
                value[i] = xssClean(value[i]);
                LogUtil.njcnDebug(log, "使用AntiSamy进行过滤清理，过滤清理之后的数据：{}", value[i]);
            }
        });
        return requestMap;
    }

    /**
     * 重写处理json数据的方法
     */
    @Override
    public ServletInputStream getInputStream() throws IOException {
        ServletInputStream inputStream = super.getInputStream();
        if (Objects.isNull(inputStream)) {
            return null;
        }
        // 读取流
        BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, StandardCharsets.UTF_8));
        // 获取json格式的数据
        ObjectMapper objectMapper = new ObjectMapper();
        String finishJson = null;
        StringBuilder sb = new StringBuilder();
        String inputStr;
        while ((inputStr = reader.readLine()) != null) {
            sb.append(inputStr);
        }
        LogUtil.njcnDebug(log, "json过滤前:{}", sb.toString());
        if (sb.toString().startsWith("[")) {
            //jsonArray
            List params = objectMapper.readValue(sb.toString(), List.class);
            int size = params.size();
            Object[] objects = new Object[params.size()];
            for (int i = 0; i < params.size(); i++) {
                Object temp = params.get(i);
                if (Objects.isNull(temp)) {
                    objects[i] = "";
                } else if (temp instanceof Number || temp instanceof List) {
//                    objects[i] = xssClean(Objects.isNull(temp) ? "" : temp.toString());
                    objects[i] = temp;
                } else {
                    objects[i] = temp;
                }
            }
            finishJson = objectMapper.writeValueAsString(objects);
        } else {
            //jsonObject
            Map map = objectMapper.readValue(sb.toString(), Map.class);
            map.keySet().forEach(k -> {
                //过滤白名单参数
                if (!ArrayUtils.contains(WHITE_PARAMETER_NAME, k)) {
                    Object temp = map.get(k);
                    if (Objects.isNull(temp)) {
                        map.put(k, temp);
                        //不递归所有数据
                    } else if (temp instanceof Number || temp instanceof List || temp.toString().startsWith("{")) {
                        map.put(k, temp);
                    } else {
                        map.put(k, xssClean(Objects.isNull(map.get(k)) ? "" : map.get(k).toString()));
                    }
                }
            });
            finishJson = objectMapper.writeValueAsString(map);
        }
        // 过滤后
        LogUtil.njcnDebug(log, "json过滤后:{}", finishJson);
        // 把json数据转为流的格式进行返回
        ByteArrayInputStream bais = new ByteArrayInputStream(finishJson.getBytes());
        return new ServletInputStream() {
            @Override
            public boolean isFinished() {
                return false;
            }

            @Override
            public boolean isReady() {
                return false;
            }

            @Override
            public void setReadListener(ReadListener listener) {
            }

            @Override
            public int read() {
                return bais.read();
            }
        };
    }

}
