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.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.util.Collections;
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;

    @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();
            }

            Optional<User> optUser = userRepository.findOneByLogin(userName);
            if (optUser.isPresent()) {
                List<Role> roles = optUser.get().getRoles().stream().collect(Collectors.toList());
                additionalInfo.put("roles", roles);
            }

            if(currentUserAuthorities.contains(AuthoritiesConstants.PARTICIPANT)) {
                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 {
    }
}

@@ -411,16 +411,75 @@
 owired;%0A
+import org.springframework.beans.factory.annotation.Value;%0A
 import o
@@ -912,24 +912,20 @@
 ava.
-u
 ti
-l.Collections
+me.Instant
 ;%0Aim
@@ -1278,24 +1278,94 @@
 epository;%0A%0A
+    @Value(%22$%7Bspring.application.name%7D%22)%0A    private String appName;%0A%0A
     @Overrid
@@ -2181,32 +2181,190 @@
 %0A            %7D%0A%0A
+            if(userName!=null) %7B%0A                // add the 'sub' claim in accordance with JWT spec%0A                additionalInfo.put(%22sub%22, userName);%0A%0A    
             Opti
@@ -2433,24 +2433,28 @@
             
+  
+  
 if (optUser.
@@ -2464,24 +2464,28 @@
 resent()) %7B%0A
+    
             
@@ -2579,32 +2579,36 @@
                 
+    
 additionalInfo.p
@@ -2642,134 +2642,16 @@
     
- %7D%0A%0A
      
-       if(currentUserAuthorities.contains(AuthoritiesConstants.PARTICIPANT)) %7B%0A                if(userName!=null) %7B%0A 
+%7D%0A%0A
     
@@ -2658,27 +2658,24 @@
             
-   
 List%3CSource%3E
@@ -2731,20 +2731,16 @@
         
-    
 .findSou
@@ -2782,28 +2782,24 @@
             
- 
-   
 List%3CString%3E
@@ -2836,20 +2836,16 @@
 tream()%0A
-    
         
@@ -2920,36 +2920,32 @@
                 
-    
 additionalInfo.p
@@ -2974,40 +2974,191 @@
 s);%0A
-%0A                %7D%0A            %7D
+            %7D%0A            // add iat and iss optional JWT claims%0A            additionalInfo.put(%22iat%22, Instant.now().getEpochSecond());%0A            additionalInfo.put(%22iss%22, appName);
 %0A   

Add sub, iat, iss claims to JWT