package top.imwork.commons.core.xss;

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

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;

public class XssHttpServletRequestWrapper extends HttpServletRequestWrapper {
    public XssHttpServletRequestWrapper(HttpServletRequest request) {
        super(request);
    }

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

    @Override
    public String getParameter(String parameter) {
        String value = super.getParameter(parameter);
        if (value == null)
            return null;
        return cleanXSS(value);
    }

    @Override
    public String getHeader(String name) {
        String value = super.getHeader(name);
        if (value == null)
            return null;
        return cleanXSS(value);
    }

    @Override
    public ServletInputStream getInputStream() throws IOException {
        final ByteArrayInputStream bai = new ByteArrayInputStream(inputHandlers(super.getInputStream()).getBytes());
        return new ServletInputStream() {
            public int read(){
                return bai.read();
            }

            public boolean isFinished() {
                return false;
            }

            public boolean isReady() {
                return false;
            }

            public void setReadListener(ReadListener readListener) {
                // Do nothing because of X and Y.
            }
        };
    }

    public String inputHandlers(ServletInputStream servletInputStream) throws IOException {
        StringBuilder sb = new StringBuilder();

        try(BufferedReader reader = new BufferedReader(new InputStreamReader(servletInputStream, StandardCharsets.UTF_8))){
            String line;
            while ((line = reader.readLine()) != null)
                sb.append(line);
        }

        return cleanXSS(sb.toString());
    }

    private String cleanXSS(String value) {
        value = value.replace("<", "&lt;").replace(">", "&gt;");
        value = value.replace("\\(", "&#40;").replace("\\)", "&#41;");
        value = value.replace("'", "&#39;");
        value = value.replace("eval\\((.*)\\)", "");
        value = value.replace("[\\\"\\'][\\s]*javascript:(.*)[\\\"\\']", "\"\"");
        String[] values = value.split(" ");
        String badStr = "'|and|exec|execute|insert|select|delete|update|count|drop|%|chr|mid|master|truncate|char|declare|sitename|net user|xp_cmdshell|;|or|-|+|,|like'|and|exec|execute|insert|create|drop|table|from|grant|use|group_concat|column_name|information_schema.columns|table_schema|union|where|select|delete|update|order|by|count|chr|mid|master|truncate|char|declare|or|;|-|--|,|like|//|/|%|#";
        String[] badStrs = badStr.split("\\|");
        for (String str : badStrs) {
            for (int k = 0; k < values.length; k++) {
                if (values[k].equalsIgnoreCase(str))
                    values[k] = "forbid";
            }
        }
        StringBuilder sb = new StringBuilder();
        for (int j = 0; j < values.length; j++) {
            if (j == values.length - 1) {
                sb.append(values[j]);
            } else {
                sb.append(values[j]).append(" ");
            }
        }
        value = sb.toString();
        return value;
    }
}
