uoa-repository-manager-service/src/main/java/eu/dnetlib/repo/manager/service/security/AuthoritiesUpdater.java

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

import org.slf4j.Logger;
import org.mitre.openid.connect.model.OIDCAuthenticationToken;
import org.slf4j.LoggerFactory;
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.FindByIndexNameSessionRepository;
import org.springframework.session.Session;
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 = LoggerFactory.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) {
                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;
                            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);
    }
}