package org.radarcns.management.security;

import org.radarcns.management.domain.Role;
import org.radarcns.management.domain.Source;
import org.radarcns.management.domain.User;
import org.radarcns.management.repository.SubjectRepository;
import org.radarcns.management.repository.UserRepository;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.security.core.authority.AuthorityUtils;
import org.springframework.security.core.userdetails.UserDetails;
import org.springframework.security.oauth2.common.DefaultOAuth2AccessToken;
import org.springframework.security.oauth2.common.OAuth2AccessToken;
import org.springframework.security.oauth2.provider.OAuth2Authentication;
import org.springframework.security.oauth2.provider.token.TokenEnhancer;

import java.time.Instant;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.stream.Collectors;

public class ClaimsTokenEnhancer implements TokenEnhancer, InitializingBean {

    @Autowired
    private SubjectRepository subjectRepository;

    @Autowired
    private UserRepository userRepository;

    @Value("${spring.application.name}")
    private String appName;

    @Override
    public OAuth2AccessToken enhance(OAuth2AccessToken accessToken,
            OAuth2Authentication authentication) {
        if (authentication.getAuthorities() != null && !authentication.getAuthorities().isEmpty()) {
            Map<String, Object> additionalInfo = new HashMap<>();
            Set<String> currentUserAuthorities = AuthorityUtils.authorityListToSet(authentication.getAuthorities());

            String userName = null;
            if (authentication.getPrincipal() instanceof UserDetails) {
                UserDetails springSecurityUser = (UserDetails) authentication.getPrincipal();
                userName = springSecurityUser.getUsername();
            } else if (authentication.getPrincipal() instanceof String) {
                userName = (String) authentication.getPrincipal();
            }

            if(userName!=null) {
                // add the 'sub' claim in accordance with JWT spec
                additionalInfo.put("sub", userName);

                Optional<User> optUser = userRepository.findOneByLogin(userName);
                if (optUser.isPresent()) {
                    List<String> roles = optUser.get().getRoles().stream().map(role ->
                        (role.getProject() == null ? "" : role.getProject().getProjectName()) +
                            ":" + role.getAuthority().getName()).collect(Collectors.toList());
                    additionalInfo.put("roles", roles);
                }

                List<Source> assignedSources = subjectRepository
                    .findSourcesBySubjectLogin(userName);

                List<String> sourceIds = assignedSources.stream()
                    .map(Source::getSourceId).collect(Collectors.toList());
                additionalInfo.put("sources", sourceIds);
            }
            // add iat and iss optional JWT claims
            additionalInfo.put("iat", Instant.now().getEpochSecond());
            additionalInfo.put("iss", appName);
            ((DefaultOAuth2AccessToken) accessToken)
                .setAdditionalInformation(additionalInfo);
        }
        return accessToken;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
    }
}

@@ -2546,20 +2546,8 @@
 am()
-.map(role -%3E
 %0A   
@@ -2567,17 +2567,32 @@
         
-(
+.filter(role -%3E 
 role.get
@@ -2605,22 +2605,53 @@
 t() 
-=
+!
 = null
- ? %22%22 :
+)%0A                        .map(role -%3E
  rol
@@ -2685,11 +2685,14 @@
 me()
-)
  +
+ %22:%22
 %0A   
@@ -2715,20 +2715,16 @@
         
- %22:%22
  + role.
@@ -2748,16 +2748,41 @@
 tName())
+%0A                        
 .collect

Remove system roles from roles claim.