package com.securitycontrol.common.security.utils;

import lombok.extern.slf4j.Slf4j;

import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.UnsupportedEncodingException;
import java.net.URLDecoder;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

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

/**
 * @author GeYazhong
 * @date 2021/11/23 10:31
 */
@Slf4j
public class XssRequestWrapper extends HttpServletRequestWrapper {
    private byte[] body;

    private String queryString;

    private String streamParam;

    private boolean checked;

    public XssRequestWrapper(HttpServletRequest request) {
        super(request);
        getParameterMap();
        BufferedReader reader;
        try {
            reader = request.getReader();
            StringBuilder sb = new StringBuilder();
            char[] buf = new char[1024];
            int rd;
            while ((rd = reader.read(buf)) != -1) {
                sb.append(buf, 0, rd);
            }
            reader.close();

            streamParam = xssClean(sb.toString());

            setChecked(xssCleanNew(sb.toString()) && xssCleanNew(request.getQueryString()));
            body = streamParam.getBytes();
        } catch (IOException e) {
            log.error(e.getLocalizedMessage(),e);
        }

        queryString = xssClean(request.getQueryString());

    }

    @Override
    public String getQueryString() {
        return queryString;
    }

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

    @Override
    public ServletInputStream getInputStream() throws IOException {

        final ByteArrayInputStream bais = new ByteArrayInputStream(body);
        return new ServletInputStream() {
            @Override
            public boolean isFinished() {
                return false;
            }

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

            @Override
            public void setReadListener(ReadListener listener) {

            }

            @Override
            public int read() throws IOException {
                int read = bais.read();
                bais.close();
                return read;
            }
        };
    }
    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> requestMap = super.getParameterMap();
        if(requestMap!=null){
            Iterator iterator = requestMap.entrySet().iterator();
            while (iterator.hasNext()) {
                Map.Entry me = (Map.Entry) iterator.next();
                String[] values = (String[]) me.getValue();
                for (int i = 0; i < values.length; i++) {
                    if (values[i] != null) {
                        values[i] = xssClean(values[i]);
                    }
                }
            }}
        return requestMap;
    }

    @Override
    public String[] getParameterValues(String paramString) {
        String[] arrayOfString1 = super.getParameterValues(paramString);
        if (arrayOfString1 == null){
            return null;
        }
        int i = arrayOfString1.length;
        String[] arrayOfString2 = new String[i];

        for (int j = 0; j < i; j++) {
            if (arrayOfString1[j] != null) {
                arrayOfString2[j] = xssClean(arrayOfString1[j]);
            }
        }
        return arrayOfString2;
    }

    @Override
    public String getParameter(String paramString) {
        String str = super.getParameter(paramString);
        if (str == null){
            return null;
        }
        return xssClean(str);
    }

    @Override
    public String getHeader(String paramString) {
        String str = super.getHeader(paramString);
        if (str == null){
            return null;
        }
        return xssClean(str);
    }
    String regex1="<script>(.*?)</script>";
    String regex2="src[\r\n]*=[\r\n]*'(.*?)'";
    String regex3="src[\r\n]*=[\r\n]*\"(.*?)\"";
    String regex4="</script>";
    String regex5="<script(.*?)>";
    String regex6="eval\\((.*?)\\)";
    String regex7="e-xpression\\((.*?)\\)";
    String regex8="javascript:";
    String regex9="vbscript:";
    String regex10="onload(.*?)=";
    private String xssClean(String value) {
        if (value == null) {
            return value;
        }
        value = value.replaceAll("", "");

        Pattern scriptPattern = Pattern.compile(regex1, Pattern.CASE_INSENSITIVE);
        value = scriptPattern.matcher(value).replaceAll("");
        scriptPattern = Pattern.compile(regex2, Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL);
        value = scriptPattern.matcher(value).replaceAll("");
        scriptPattern = Pattern.compile(regex3, Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL);
        value = scriptPattern.matcher(value).replaceAll("");
        scriptPattern = Pattern.compile(regex4, Pattern.CASE_INSENSITIVE);
        value = scriptPattern.matcher(value).replaceAll("");
        scriptPattern = Pattern.compile(regex5, Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL);
        value = scriptPattern.matcher(value).replaceAll("");
        scriptPattern = Pattern.compile(regex6, Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL);
        value = scriptPattern.matcher(value).replaceAll("");
        scriptPattern = Pattern.compile(regex7, Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL);
        value = scriptPattern.matcher(value).replaceAll("");
        scriptPattern = Pattern.compile(regex8, Pattern.CASE_INSENSITIVE);
        value = scriptPattern.matcher(value).replaceAll("");
        scriptPattern = Pattern.compile(regex9, Pattern.CASE_INSENSITIVE);
        value = scriptPattern.matcher(value).replaceAll("");
        scriptPattern = Pattern.compile(regex10, Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL);
        value = scriptPattern.matcher(value).replaceAll("");

        return value;
    }


    private boolean xssCleanNew(String value) {
        boolean find = false;
        if (value == null) {
            return true;
        }
        List<Pattern> patterns = new ArrayList<Pattern>();
        patterns.add(Pattern.compile(regex1, Pattern.CASE_INSENSITIVE));
        patterns.add(Pattern.compile(regex2, Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL));
        patterns.add(Pattern.compile(regex3, Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL));
        patterns.add(Pattern.compile(regex4, Pattern.CASE_INSENSITIVE));
        patterns.add(Pattern.compile(regex5, Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL));
        patterns.add(Pattern.compile(regex6, Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL));
        patterns.add(Pattern.compile(regex7, Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL));
        patterns.add(Pattern.compile(regex8, Pattern.CASE_INSENSITIVE));
        patterns.add(Pattern.compile(regex9, Pattern.CASE_INSENSITIVE));
        patterns.add(Pattern.compile(regex10, Pattern.CASE_INSENSITIVE | Pattern.MULTILINE | Pattern.DOTALL));

        for (Pattern pattern : patterns) {
            find = match(pattern, value) ;
            if (find) {
                break;
            }

        }
        return !find;
    }

    /**
     * 执行正则表达式
     *
     * @param pattern
     *            表达式
     * @param str
     *            待验证字符串
     * @return 返回 <b>true </b>,否则为 <b>false </b>
     */
    private static boolean match(Pattern pattern, String str) {
        Matcher m = pattern.matcher(str);
        return m.find();
    }

    public String getReaderParam() {
        return streamParam;
    }

    public boolean isChecked() {
        return checked;
    }

    public void setChecked(boolean checked) {
        this.checked = checked;
    }
}
