package com.atlassian.bitbucket.security;

import com.atlassian.stash.internal.user.StashUserAuthenticationToken;
import com.google.common.base.Defaults;
import com.google.common.collect.Lists;
import org.apache.commons.lang3.ClassUtils;
import org.junit.runners.model.MultipleFailureException;
import org.mockito.invocation.InvocationOnMock;
import org.mockito.stubbing.Answer;
import org.springframework.context.ApplicationContext;
import org.springframework.context.support.ClassPathXmlApplicationContext;
import org.springframework.core.LocalVariableTableParameterNameDiscoverer;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.security.access.AccessDeniedException;
import org.springframework.security.access.expression.method.DefaultMethodSecurityExpressionHandler;
import org.springframework.security.access.prepost.PostAuthorize;
import org.springframework.security.access.prepost.PreAuthorize;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.util.ReflectionUtils;

import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.util.List;

import static org.junit.Assert.fail;
import static org.mockito.Mockito.*;

/**
 * Ensures all {@code PreAuthorize} and {@code PostAuthorize} annotations are <i>somewhat</i> correct, which is to say
 * they reference <i>real</i> fields and <i>real</i> parameters. Otherwise mistakes of this kind are ignored and the
 * permission check is not applied, which could be disastrous for security.
 */
public class AuthorizedAnnotationHelper {

    private final Class<?> serviceClass;

    public AuthorizedAnnotationHelper(Class<?> serviceClass) {
        this.serviceClass = serviceClass;
    }

    public void testAuthorizedMethods() throws Throwable {
        MultipleFailureException.assertEmpty(findBrokenAuthorizedMethods());
    }

    public List<Throwable> findBrokenAuthorizedMethods() throws Exception {
        // Use a null user so isAuthenticated() fails
        SecurityContextHolder.getContext().setAuthentication(new StashUserAuthenticationToken.Builder().build());
        try {
            return validateAnnotations(serviceClass);
        } finally {
            SecurityContextHolder.getContext().setAuthentication(null);
        }
    }

    private static Object[] mockArguments(Method method) {
        Class<?>[] parameterTypes = method.getParameterTypes();
        // We just need to make sure that are arguments are never null and return non-null objects from some properties
        Object[] args = new Object[parameterTypes.length];
        for (int i = 0; i < parameterTypes.length; i++) {
            // Rely on Mockito for everything except for what it can't handle (ie final classes)
            if (parameterTypes[i].isPrimitive()) {
                args[i] = Defaults.defaultValue(parameterTypes[i]);
            } else if (parameterTypes[i] == String.class) {
                // Make sure this isn't null for MockMethodSecurityExpressionHandler
                args[i] = "";
            } else if (Modifier.isFinal(parameterTypes[i].getModifiers())) {
                args[i] = null;
            } else {
                // Deep mock so nested EL expressions (ie a.b) never return null
                args[i] = mock(parameterTypes[i], RETURNS_DEEP_STUBS);
            }
        }
        return args;
    }

    // Wrap bean in Spring security to activate Authorize annotations
    private static Object secureBean(final Class<?> serviceClass, Object bean) {
        ApplicationContext securityContext = new ClassPathXmlApplicationContext("classpath:/spring/security-test-context.xml");
        securityContext.getBean(DefaultMethodSecurityExpressionHandler.class)
                .setParameterNameDiscoverer(new LocalVariableTableParameterNameDiscoverer() {
                    @Override
                    public String[] getParameterNames(Method method) {
                        // Because we're spying we need to get parameter names from the _actual_ class
                        method = ReflectionUtils.findMethod(serviceClass, method.getName(), method.getParameterTypes());
                        return super.getParameterNames(method);
                    }
                });
        return securityContext.getAutowireCapableBeanFactory().initializeBean(bean, serviceClass.getName());
    }

    private static List<Throwable> validateAnnotations(Class<?> serviceClass) throws Exception {
        // Check for a class-level annotation. If it's present, it means every visible method needs to be checked,
        // not just methods that are explicitly annotated
        //
        // Note: PostAuthorize is not checked at the class level because there is nowhere in the codebase where it
        // is applied that way
        PreAuthorize preAuthorize = AnnotationUtils.findAnnotation(serviceClass, PreAuthorize.class);

        // Create a service by bypassing the constructor - we don't actually care about any of the fields
        // If the @PreAuthorize fails then we expect a NPE at some stage (if the method is non-trivial)
        Object spy = mock(serviceClass);
        Object secureSpy = secureBean(serviceClass, spy);
        List<Throwable> exceptions = Lists.newArrayList();
        // Ideally we get the declared methods from the class, but we can't because it gets mocked
        for (Class<?> serviceInterface : ClassUtils.getAllInterfaces(serviceClass)) {
            for (Method serviceMethod : serviceInterface.getDeclaredMethods()) {
                Method method = ReflectionUtils.findMethod(serviceClass, serviceMethod.getName(), serviceMethod.getParameterTypes());
                try {
                    // TODO We need to parse the value of both of these types and 'test' all of the conditional paths
                    // e.g. hasGlobalPermission(#good) and hasGlobalPermission(#bad)
                    // This will pass even though 'bad' may not exist
                    if (AnnotationUtils.findAnnotation(method, PostAuthorize.class) != null) {
                        validatePostAuthorize(spy, secureSpy, serviceMethod, mockArguments(method));
                    } else if (preAuthorize != null || AnnotationUtils.findAnnotation(method, PreAuthorize.class) != null) {
                        validatePreAuthorize(secureSpy, serviceMethod, mockArguments(method));
                    }
                } catch (AccessDeniedException e) {
                    // Expected
                } catch (Throwable e) {
                    // Test all the methods instead of failing on the first one
                    exceptions.add(e);
                }
            }
        }
        return exceptions;
    }

    private static void validatePreAuthorize(Object secureSpy, Method serviceMethod, Object[] args) {
        // Invoke the method with mock arguments
        ReflectionUtils.invokeMethod(serviceMethod, secureSpy, args);
        fail("PreAuthorize has failed: " + serviceMethod);
    }

    private static void validatePostAuthorize(Object spy, Object secureSpy, Method serviceMethod, Object[] args)
            throws Exception {
        // Return something from this method when is then checked by PostAuthorize
        serviceMethod.invoke(doAnswer(new Answer() {
            @Override
            public Object answer(InvocationOnMock invocation) throws Throwable {
                // Null is fine 99% of the time, but we might get a property from the return type
                return mock(invocation.getMethod().getReturnType(), RETURNS_DEEP_STUBS);
            }
        }).when(spy), args);
        ReflectionUtils.invokeMethod(serviceMethod, secureSpy, args);
        fail("PostAuthorize has failed: " + serviceMethod);
    }
}
