package cn.myapps.runtime.security;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.*;

public class PayloadXssRequestWrapper extends HttpServletRequestWrapper {



    private final String body;

    /**
     * Constructs a request object wrapping the given request.
     *
     * @param request
     * @throws IllegalArgumentException if the request is null
     */
    public PayloadXssRequestWrapper(HttpServletRequest request) {
        super(request);
        StringBuilder stringBuilder = new StringBuilder();
        BufferedReader bufferedReader = null;
        try {
            InputStream inputStream = request.getInputStream();
            if (inputStream != null) {
                bufferedReader = new BufferedReader(new InputStreamReader(inputStream));
                char[] charBuffer = new char[1024];
                int bytesRead = -1;
                while ((bytesRead = bufferedReader.read(charBuffer)) > 0) {
                    stringBuilder.append(charBuffer, 0, bytesRead);
                }
            } else {
                stringBuilder.append("");
            }
        } catch (IOException ex) {
            ex.printStackTrace();
        } finally {
            if (bufferedReader != null) {
                try {
                    bufferedReader.close();
                } catch (IOException ex) {
                    ex.printStackTrace();
                }
            }
        }
        body = stringBuilder.toString();
    }


    //替换非法字符
    private String clean(String s) {
        boolean isRuntimeURI = getRequestURI().contains("/api/runtime");
        StringBuilder sb = new StringBuilder(s.length() + 16);
        for (int i = 0; i < s.length(); i++)
        {
            char c = s.charAt(i);
            switch (c)
            {
                case '>':
                    sb.append('＞');// 全角大于号
                    break;
                case '<':
                    sb.append('＜');// 全角小于号
                    break;
                case '\'':
                    if(isRuntimeURI){
                        sb.append('‘');// 全角单引号
                    } else {
                        sb.append('\'');
                    }
                    break;
//            case '\"':
//                sb.append('“');// 全角双引号
//                break;
//            case '\\':
//                sb.append('＼');// 全角斜线
//                break;
                case '%':
                    sb.append('％'); // 全角冒号
                    break;
                case '(':
                    sb.append('（'); // 全角左括号
                    break;
                case ')':
                    sb.append('）'); // 全角右边括号
                    break;
                default:
                    sb.append(c);
                    break;
            }
        }
        return sb.toString();
    }




    /**
     * Read Request Body payload  in Filter
     * 读取输入流,,文件post时采用这种方式
     *
     * @return
     * @throws IOException
     */
    @Override
    public BufferedReader getReader() throws IOException {
        return new BufferedReader(new InputStreamReader(this.getInputStream()));
    }


    /**
     * Read Request Body payload  in Filter
     * 读取输入流,,文件post时采用这种方式
     *
     * @return
     * @throws IOException
     */
    @Override
    public ServletInputStream getInputStream() throws IOException {
        System.out.println("getInputStream");
        final ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(clean(body).getBytes());//在这里过滤非法字符
        ServletInputStream servletInputStream = new ServletInputStream() {
            public boolean isFinished() {
                return false;
            }

            public boolean isReady() {
                return false;
            }

            public int read() throws IOException {
                return byteArrayInputStream.read();
            }

            @Override
            public void setReadListener(ReadListener readListener) {
                // TODO Auto-generated method stub

            }
        };
        return servletInputStream;
    }

}
