package com.tcm.training.admin.common.xss;

import org.springframework.core.MethodParameter;
import org.springframework.http.HttpInputMessage;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.RequestBodyAdviceAdapter;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Field;
import java.lang.reflect.Modifier;
import java.lang.reflect.Type;
import java.util.Collection;
import java.util.IdentityHashMap;
import java.util.Map;
import java.util.Set;

@ControllerAdvice(annotations = Controller.class)
public class XssRequestBodyAdvice extends RequestBodyAdviceAdapter {

    @Override
    public boolean supports(MethodParameter methodParameter, Type targetType, Class<? extends HttpMessageConverter<?>> converterType) {
        return true;
    }

    @Override
    public Object afterBodyRead(Object body,
                                HttpInputMessage inputMessage,
                                MethodParameter parameter,
                                Type targetType,
                                Class<? extends HttpMessageConverter<?>> converterType) {
        if (!(inputMessage instanceof ServletServerHttpRequest)) {
            return body;
        }
        HttpServletRequest request = ((ServletServerHttpRequest) inputMessage).getServletRequest();
        String contentType = request.getContentType();
        if (contentType == null || !contentType.toLowerCase().contains("application/json")) {
            return body;
        }
        sanitize(body, new IdentityHashMap<Object, Boolean>());
        return body;
    }

    @SuppressWarnings("unchecked")
    private void sanitize(Object target, Map<Object, Boolean> visited) {
        if (target == null || visited.containsKey(target)) {
            return;
        }
        visited.put(target, Boolean.TRUE);

        if (target instanceof String || target instanceof Number || target instanceof Boolean || target instanceof Enum) {
            return;
        }
        if (target instanceof Collection) {
            for (Object item : (Collection<Object>) target) {
                sanitize(item, visited);
            }
            return;
        }
        if (target instanceof Map) {
            for (Map.Entry<Object, Object> entry : ((Map<Object, Object>) target).entrySet()) {
                Object value = entry.getValue();
                if (value instanceof String) {
                    entry.setValue(XssCleaner.cleanPlainText((String) value));
                } else {
                    sanitize(value, visited);
                }
            }
            return;
        }
        if (target.getClass().isArray()) {
            return;
        }

        Class<?> type = target.getClass();
        while (type != null && !Object.class.equals(type)) {
            Field[] fields = type.getDeclaredFields();
            for (Field field : fields) {
                if (Modifier.isStatic(field.getModifiers()) || Modifier.isFinal(field.getModifiers())) {
                    continue;
                }
                field.setAccessible(true);
                try {
                    Object value = field.get(target);
                    if (value == null) {
                        continue;
                    }
                    if (value instanceof String) {
                        String sanitized = field.isAnnotationPresent(XssSafeHtml.class)
                                ? XssCleaner.cleanRichText((String) value)
                                : XssCleaner.cleanPlainText((String) value);
                        field.set(target, sanitized);
                    } else {
                        sanitize(value, visited);
                    }
                } catch (IllegalAccessException ignored) {
                }
            }
            type = type.getSuperclass();
        }
    }
}