package org.jeecg.config.filter;

import cn.hutool.core.io.IoUtil;
import cn.hutool.core.util.StrUtil;

import lombok.extern.slf4j.Slf4j;
import org.jeecg.common.constant.SecurityFilterProperties;
import org.jeecg.common.exception.JeecgBootException;
import org.jeecg.common.util.SecurityUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletRequestWrapper;
import java.io.BufferedReader;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * XSS 和 SQL 注入防护过滤器
 * <p>
 * 对所有 HTTP 请求进行安全检查，包括：
 * <ul>
 *     <li>请求参数（Query String 和 Form Data）</li>
 *     <li>请求体（JSON、XML 等）</li>
 *     <li>请求头</li>
 * </ul>
 *
 * @author zsElectric
 */
@Slf4j
public class XssAndSqlInjectionFilter implements Filter {

    private final SecurityFilterProperties properties;

    /**
     * 默认排除的 URL 路径（不进行安全检查）
     */
    private static final Set<String> DEFAULT_EXCLUDE_URLS = new HashSet<>(Arrays.asList(
            "/api/v1/auth/captcha",
            "/doc.html",
            "/swagger-ui",
            "/v3/api-docs",
            "/webjars"
    ));

    /**
     * 默认需要检查的请求头
     */
    private static final Set<String> DEFAULT_CHECK_HEADERS = new HashSet<>(Arrays.asList(
            "Referer",
            "X-Forwarded-For"
    ));

    public XssAndSqlInjectionFilter(SecurityFilterProperties properties) {
        this.properties = properties;
    }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain)
            throws IOException, ServletException {

        HttpServletRequest httpRequest = (HttpServletRequest) request;

        // 检查是否为排除的 URL
        if (isExcludedUrl(httpRequest)) {
            chain.doFilter(request, response);
            return;
        }

        try {
            // 包装请求，进行安全检查
            SecurityRequestWrapper wrappedRequest = new SecurityRequestWrapper(httpRequest);
            chain.doFilter(wrappedRequest, response);
        } catch (Exception e) {
            log.error("检测到恶意请求，URL: {}, 错误信息: {}", httpRequest.getRequestURI(), e.getMessage());
            throw new JeecgBootException("用户输入包含非法内容，请输入合法内容!");
        }
    }

    /**
     * 检查 URL 是否在排除列表中
     */
    private boolean isExcludedUrl(HttpServletRequest request) {
        String uri = request.getRequestURI();
        
        // 检查默认排除列表
        if (DEFAULT_EXCLUDE_URLS.stream().anyMatch(uri::startsWith)) {
            return true;
        }
        
        // 检查配置的排除列表
        if (properties.getExcludeUrls() != null) {
            return properties.getExcludeUrls().stream().anyMatch(uri::startsWith);
        }
        
        return false;
    }

    /**
     * 安全请求包装类
     */
    private class SecurityRequestWrapper extends HttpServletRequestWrapper {

        private byte[] body;

        public SecurityRequestWrapper(HttpServletRequest request) throws IOException {
            super(request);
            // 缓存请求体
            if (isJsonOrXmlRequest(request)) {
                body = IoUtil.readBytes(request.getInputStream());
                // 检查请求体
                String bodyContent = new String(body, StandardCharsets.UTF_8);
                if (StrUtil.isNotBlank(bodyContent)) {
                    checkContent(bodyContent, "请求体");
                }
            }
            // 检查请求头
            checkHeaders();
        }

        @Override
        public String getParameter(String name) {
            String value = super.getParameter(name);
            if (value != null) {
                checkContent(value, "请求参数[" + name + "]");
            }
            return value;
        }

        @Override
        public String[] getParameterValues(String name) {
            String[] values = super.getParameterValues(name);
            if (values != null) {
                for (int i = 0; i < values.length; i++) {
                    if (values[i] != null) {
                        checkContent(values[i], "请求参数[" + name + "][" + i + "]");
                    }
                }
            }
            return values;
        }

        @Override
        public Map<String, String[]> getParameterMap() {
            Map<String, String[]> parameterMap = super.getParameterMap();
            if (parameterMap != null) {
                for (Map.Entry<String, String[]> entry : parameterMap.entrySet()) {
                    String[] values = entry.getValue();
                    if (values != null) {
                        for (int i = 0; i < values.length; i++) {
                            if (values[i] != null) {
                                checkContent(values[i], "请求参数[" + entry.getKey() + "][" + i + "]");
                            }
                        }
                    }
                }
            }
            return parameterMap;
        }

        @Override
        public ServletInputStream getInputStream() throws IOException {
            if (body == null) {
                return super.getInputStream();
            }

            final ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(body);

            return new ServletInputStream() {
                @Override
                public boolean isFinished() {
                    return byteArrayInputStream.available() == 0;
                }

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

                @Override
                public void setReadListener(ReadListener readListener) {
                    // Not implemented
                }

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

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

        /**
         * 检查请求头
         */
        private void checkHeaders() {
            Set<String> headersToCheck = new HashSet<>(DEFAULT_CHECK_HEADERS);
            if (properties.getCheckHeaders() != null && !properties.getCheckHeaders().isEmpty()) {
                headersToCheck.addAll(properties.getCheckHeaders());
            }
            
            for (String headerName : headersToCheck) {
                String headerValue = super.getHeader(headerName);
                if (headerValue != null) {
                    checkContent(headerValue, "请求头[" + headerName + "]");
                }
            }
        }

        /**
         * 检查内容是否包含恶意代码
         *
         * @param content 待检查的内容
         * @param location 内容来源位置（用于日志记录）
         */
        private void checkContent(String content, String location) {
            // XSS 检测
            if (properties.getXssEnabled() && SecurityUtils.containsXss(content)) {
//                throw new JeecgBootException("检测到 XSS 攻击尝试，位置: " + location);
                throw new JeecgBootException("用户输入包含非法内容，请输入合法内容!");
            }

            // SQL 注入检测
            if (properties.getSqlInjectionEnabled() && SecurityUtils.containsSqlInjection(content)) {
                throw new JeecgBootException("用户输入包含非法内容，请输入合法内容!");
            }
        }

        /**
         * 判断是否为 JSON 或 XML 请求
         */
        private boolean isJsonOrXmlRequest(HttpServletRequest request) {
            String contentType = request.getHeader(HttpHeaders.CONTENT_TYPE);
            if (contentType == null) {
                return false;
            }
            return contentType.contains(MediaType.APPLICATION_JSON_VALUE)
                    || contentType.contains(MediaType.APPLICATION_XML_VALUE)
                    || contentType.contains(MediaType.TEXT_XML_VALUE);
        }
    }

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
        log.info("XSS 和 SQL 注入防护过滤器已启动");
    }

    @Override
    public void destroy() {
        log.info("XSS 和 SQL 注入防护过滤器已销毁");
    }
}
