D-Net

package eu.dnetlib.repo.manager.service.security;

import org.apache.log4j.Logger;

import org.mitre.openid.connect.model.OIDCAuthenticationToken;

import org.springframework.beans.factory.annotation.Autowired;

import org.springframework.security.core.Authentication;

import org.springframework.security.core.GrantedAuthority;

import org.springframework.security.core.context.SecurityContext;

import org.springframework.security.core.context.SecurityContextHolder;

import org.springframework.security.oauth2.common.exceptions.UnauthorizedClientException;

import org.springframework.security.web.context.HttpSessionSecurityContextRepository;

import org.springframework.session.ExpiringSession;

import org.springframework.session.FindByIndexNameSessionRepository;

import org.springframework.stereotype.Service;

import java.util.Collection;

import java.util.HashSet;

import java.util.Map;

@Service

public class AuthoritiesUpdater extends HttpSessionSecurityContextRepository {

    private static final Logger logger = Logger.getLogger(AuthoritiesUpdater.class);

    @Autowired

    FindByIndexNameSessionRepository sessions;

    public void update(String id, Collection<? extends GrantedAuthority> authorities) {

        if (sessions != null) {

            Map<String, ExpiringSession> map = sessions.

                    findByIndexNameAndIndexValue(FindByIndexNameSessionRepository.PRINCIPAL_NAME_INDEX_NAME, id);

            logger.debug(map.values().toArray().length);

            for (ExpiringSession session : map.values()) {

                logger.debug(session.getId());

                if (!session.isExpired()) {

                    SecurityContext securityContext = session.getAttribute(SPRING_SECURITY_CONTEXT_KEY);

                    Authentication authentication = securityContext.getAuthentication();

                    if (authentication instanceof OIDCAuthenticationToken) {

                        OIDCAuthenticationToken authOIDC = (OIDCAuthenticationToken) authentication;

                        logger.debug(authorities);

                        securityContext.setAuthentication(new OIDCAuthenticationToken(authOIDC.getSub(), authOIDC.getIssuer(),

                                authOIDC.getUserInfo(), authorities, authOIDC.getIdToken(),

                                authOIDC.getAccessTokenValue(), authOIDC.getRefreshTokenValue()));

                        logger.debug("Update authorities");

                        session.setAttribute(SPRING_SECURITY_CONTEXT_KEY, securityContext);

                        sessions.save(session);

                    }

                }

            }

        }

    }

    public void update(String id, Update update) {

        Collection<? extends GrantedAuthority> authorities = update.authorities(SecurityContextHolder.getContext().getAuthentication().getAuthorities());

        this.update(id, authorities);

    }

    public void addRole(GrantedAuthority role) {

        Authentication auth = SecurityContextHolder.getContext().getAuthentication();

        if (auth instanceof OIDCAuthenticationToken) {

            OIDCAuthenticationToken oidcAuth = (OIDCAuthenticationToken) auth;

            this.update(oidcAuth.getUserInfo().getEmail(), old -> {

                HashSet<GrantedAuthority> authorities = new HashSet<>(old);

                authorities.add(role);

                return authorities;

            });

        } else {

            throw new UnauthorizedClientException("User auth is not instance of OIDCAuthenticationToken");

        }

    }

    public void removeRole(GrantedAuthority role) {

        Authentication auth = SecurityContextHolder.getContext().getAuthentication();

        if (auth instanceof OIDCAuthenticationToken) {

            OIDCAuthenticationToken oidcAuth = (OIDCAuthenticationToken) auth;

            this.update(oidcAuth.getUserInfo().getEmail(), old -> {

                HashSet<GrantedAuthority> authorities = new HashSet<>(old);

                authorities.remove(role);

                return authorities;

            });

        } else {

            throw new UnauthorizedClientException("User auth is not instance of OIDCAuthenticationToken");

        }

    }

    public interface Update {

        Collection<? extends GrantedAuthority> authorities(Collection<? extends GrantedAuthority> old);

    }

}