package com.saye.hrs.config;

import com.alibaba.fastjson.JSONObject;
import com.saye.hrs.commons.string.StringDUtil;
import org.jsoup.Jsoup;
import org.jsoup.safety.Whitelist;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.ByteArrayInputStream;
import org.apache.commons.io.IOUtils;
import java.io.IOException;

public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper{

	HttpServletRequest orgRequest = null;

	private static final Whitelist whitelist = Whitelist.basicWithImages();

	public XssHttpServletRequestWrapper(HttpServletRequest request) {
	    super(request);
	    orgRequest = request;
	}

	
	@Override
	public ServletInputStream getInputStream() throws IOException {
		ServletInputStream stream = null; 
		try {  
            stream = orgRequest.getInputStream();  
            String requestBody = IOUtils.toString(stream, "utf-8");
            requestBody = Jsoup.clean(requestBody,whitelist);
            JSONObject resultJson = JSONObject.parseObject(requestBody);
            byte[] array = resultJson.toString().getBytes("utf-8");
            if (array == null) {
            	array = new byte[0];
            }
            final ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(array);

            return new ServletInputStream() {
                @Override
                public boolean isFinished() {
                    return false;
                }
                @Override
                public boolean isReady() {
                    return false;
                }
                @Override
                public void setReadListener(ReadListener readListener) {
                }
                @Override
                public int read() {
                    return byteArrayInputStream.read();
                }
            };
        } catch (IOException e) {  
            e.printStackTrace();  
        }  
		return super.getInputStream();
	}


	/**
	 * 覆盖getParameter方法，将参数名和参数值都做xss过滤.
	 * 如果需要获得原始的值，则通过super.getParameterValues(name)来获取
	 * getParameterNames,getParameterValues和getParameterMap也可能需要覆盖
	 */
	@Override
	public String getParameter(String name) {
	    if (("content".equals(name) || name.endsWith("WithHtml"))) {
	        return super.getParameter(name);
	    }
	    name = Jsoup.clean(name,whitelist);
	    String value = super.getParameter(name);
	    if (StringDUtil.isNotBlank(value)) {
	        value = Jsoup.clean(value,whitelist);
	    }
	    return value;
	}
	
	@Override
	public String[] getParameterValues(String name) {
	    String[] arr = super.getParameterValues(name);
	    if (arr != null) {
	        for (int i = 0; i < arr.length; i++) {
	            arr[i] = Jsoup.clean(arr[i],whitelist);
	        }
	    }
	    return arr;
	}


	/**
	 * 覆盖getHeader方法，将参数名和参数值都做xss过滤。<br/>
	 * 如果需要获得原始的值，则通过super.getHeaders(name)来获取<br/>
	 * getHeaderNames 也可能需要覆盖
	 */
	@Override
	public String getHeader(String name) {
	    name = Jsoup.clean(name,whitelist);
	    String value = super.getHeader(name);
	    if (StringDUtil.isNotBlank(value)) {
	        value = Jsoup.clean(value,whitelist);
	    }
	    return value;
	}

	/**
	 * 获取最原始的request
	 *
	 * @return
	 */
	public HttpServletRequest getOrgRequest() {
	    return orgRequest;
	}

	/**
	 * 获取最原始的request的静态方法
	 *
	 * @return
	 */
	public static HttpServletRequest getOrgRequest(HttpServletRequest req) {
	    if (req instanceof XssHttpServletRequestWrapper) {
	        return ((XssHttpServletRequestWrapper) req).getOrgRequest();
	    }
	    return req;
	}
}