package org.radarcns.management.security;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.stream.Collectors;
import org.radarcns.management.domain.Source;
import org.radarcns.management.repository.SubjectRepository;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.beans.factory.annotation.Autowired;
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;

public class ClaimsTokenEnhancer implements TokenEnhancer, InitializingBean {

    @Autowired
    private SubjectRepository subjectRepository;

    @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());
//            additionalInfo.put("authorities" , currentUserAuthorities);
            if(currentUserAuthorities.contains(AuthoritiesConstants.PARTICIPANT)) {
                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) {
                    List<Source> assignedSources = subjectRepository
                        .findSourcesBySubjectLogin(userName);

                    List<String> sourceIds = assignedSources.stream()
                        .map(Source::getSourceId).collect(Collectors.toList());
                    additionalInfo.put("sources", sourceIds);

                }
            }
            ((DefaultOAuth2AccessToken) accessToken)
                .setAdditionalInformation(additionalInfo);
        }
        return accessToken;
    }

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

@@ -47,128 +47,43 @@
 ort 
-java.util.HashMap;%0Aimport java.util.List;%0Aimport java.util.Map;%0Aimport java.util.Set;%0Aimport java.util.stream.Collectors
+org.radarcns.management.domain.Role
 ;%0Aim
@@ -126,16 +126,60 @@
 Source;%0A
+import org.radarcns.management.domain.User;%0A
 import o
@@ -231,16 +231,74 @@
 sitory;%0A
+import org.radarcns.management.repository.UserRepository;%0A
 import o
@@ -837,16 +837,203 @@
 ancer;%0A%0A
+import java.util.Collections;%0Aimport java.util.HashMap;%0Aimport java.util.List;%0Aimport java.util.Map;%0Aimport java.util.Optional;%0Aimport java.util.Set;%0Aimport java.util.stream.Collectors;%0A%0A
 public c
@@ -1168,16 +1168,75 @@
 itory;%0A%0A
+    @Autowired%0A    private UserRepository userRepository;%0A%0A
     @Ove
@@ -1648,170 +1648,9 @@
 ));%0A
-//            additionalInfo.put(%22authorities%22 , currentUserAuthorities);%0A            if(currentUserAuthorities.contains(AuthoritiesConstants.PARTICIPANT)) %7B%0A    
+%0A
     
@@ -1689,28 +1689,24 @@
             
-  
-  
 if (authenti
@@ -1749,28 +1749,24 @@
 rDetails) %7B%0A
-    
             
@@ -1855,36 +1855,32 @@
                 
-  
-  
 userName = sprin
@@ -1904,28 +1904,24 @@
 Username();%0A
-    
             
@@ -1990,36 +1990,32 @@
                 
-   
- 
 userName = (Stri
@@ -2053,37 +2053,401 @@
 ();%0A            
-    %7D
+%7D%0A%0A            Optional%3CUser%3E optUser = userRepository.findOneByLogin(userName);%0A            if (optUser.isPresent()) %7B%0A                List%3CRole%3E roles = optUser.get().getRoles().stream().collect(Collectors.toList());%0A                additionalInfo.put(%22roles%22, roles);%0A            %7D%0A%0A            if(currentUserAuthorities.contains(AuthoritiesConstants.PARTICIPANT)) %7B
 %0A               

Add roles claim to JWT