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.Session;

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, Update update) {

        if (sessions != null) {

            Map<String, Session> map = sessions.

                    findByIndexNameAndIndexValue(FindByIndexNameSessionRepository.PRINCIPAL_NAME_INDEX_NAME, id);

            if (map != null) {

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

                for (Session 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(update.authorities(authOIDC.getAuthorities()));

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

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

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

                            logger.debug("Update authorities");

                            session.setAttribute(SPRING_SECURITY_CONTEXT_KEY, securityContext);

                            sessions.save(session);

                        }

                    }

                }

            }

        }

    }

    public void addRole(String id, GrantedAuthority role) {

        this.update(id, old -> {

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

            authorities.add(role);

            return authorities;

        });

    }

    public void addRole(GrantedAuthority role) {

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

        if (auth instanceof OIDCAuthenticationToken) {

            OIDCAuthenticationToken oidcAuth = (OIDCAuthenticationToken) auth;

            this.addRole(oidcAuth.getUserInfo().getEmail(), role);

        } else {

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

        }

    }

    public void removeRole(String id, GrantedAuthority role) {

        this.update(id, old -> {

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

            authorities.remove(role);

            return authorities;

        });

    }

    public void removeRole(GrantedAuthority role) {

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

        if (auth instanceof OIDCAuthenticationToken) {

            OIDCAuthenticationToken oidcAuth = (OIDCAuthenticationToken) auth;

            this.removeRole(oidcAuth.getUserInfo().getEmail(), role);

        }

    }

    public interface Update {

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

    }

}