package kr.co.uplus.ez.config.filter;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.charset.StandardCharsets;

import javax.servlet.ReadListener;
import javax.servlet.ServletInputStream;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;

import org.apache.poi.util.IOUtils;

import kr.co.uplus.ez.common.xss.XssPreventer;

public final class XssPreventWrapper extends HttpServletRequestWrapper {

    private byte[] bytes;
    private String requestBody;

    public XssPreventWrapper(HttpServletRequest servletRequest) throws IOException {
        super(servletRequest);

        InputStream in = super.getInputStream();
        bytes = IOUtils.toByteArray(in);
        requestBody = XssPreventer.escape(new String(bytes));
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        ByteArrayInputStream bis = new ByteArrayInputStream(this.requestBody.getBytes(StandardCharsets.UTF_8));
        ServletInputStream sis = new ServletImpl(bis);
        return sis;
    }

    public String getRequestBody() {
        return this.requestBody;
    }

    public String[] getParameterValues(String parameter) {

        String[] values = super.getParameterValues(parameter);
        if (values == null) {
            return null;
        }
        int count = values.length;
        String[] encodedValues = new String[count];
        for (int i = 0; i < count; i++) {
            encodedValues[i] = XssPreventer.escape(values[i]);
        }
        return encodedValues;
    }

    public String getParameter(String parameter) {
        String value = super.getParameter(parameter);
        if (value == null) {
            return null;
        }
        return XssPreventer.escape(value);
    }

    public String getHeader(String name) {
        String value = super.getHeader(name);
        if (value == null)
            return null;
        return XssPreventer.escape(value);
    }

    class ServletImpl extends ServletInputStream {
        private InputStream is;

        public ServletImpl(InputStream bis){
            is = bis;
        }

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

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

        @Override
        public void setReadListener(ReadListener listener) {
        }

        @Override
        public int read() throws IOException {
            return is.read();
        }

        @Override
        public int read(byte[] b) throws IOException {
            return is.read(b);
        }
    }

}
