package org.springblade.common.auth.security;

import com.alibaba.fastjson.JSON;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.text.StringEscapeUtils;
import org.springframework.util.StreamUtils;

import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.UnsupportedEncodingException;
import java.util.Map;

public class SecurityHttpRequestWrapper extends HttpServletRequestWrapper {

    private byte[] bodyCopier;

    public static ThreadLocal<String> bodyThreadLocal = new ThreadLocal<>();

    /**
     * Request body json参数转化，防XSS攻击
     */
    @SuppressWarnings("unchecked")
    public SecurityHttpRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        bodyCopier = StreamUtils.copyToByteArray(request.getInputStream());

        String body = new String(this.bodyCopier, "UTF-8");

        bodyThreadLocal.set(body);

        if (StringUtils.isBlank(body)) {
            bodyCopier = "".getBytes("UTF-8");
            return;
        }
        if (body.startsWith("[")) {//不过滤数组类型
            bodyCopier = body.getBytes("UTF-8");
            return;
        }
        Map<String, Object> map = JSON.parseObject(body, Map.class);
        StringBuilder sb = new StringBuilder();
        for (String key : map.keySet()) {
            Object value = map.get(key);
            sb.append(" key:" + key + " value:" + value);
            if (value instanceof String) {
                //map.put(key, ((String) value).toUpperCase());
                map.put(key, StringEscapeUtils.escapeHtml4((String) value));
            }
        }
        body = JSON.toJSONString(map);
        bodyCopier = body.getBytes("UTF-8");
    }

    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(this.getInputStream()));
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        return new ServletInputStreamCopier(bodyCopier);
    }

    @Override
    public String getParameter(String name) {
        String value = this.getRequest().getParameter(name);
        if (StringUtils.isBlank(value)) {
            return value;
        }
        return StringEscapeUtils.escapeHtml4((String) value);
    }

    public byte[] getCopy() {
        return this.bodyCopier;
    }

    public String getBody() throws UnsupportedEncodingException {
        return new String(this.bodyCopier, "UTF-8");
    }
}
