From b4cee1f2aa0b0b7fd5e6dc6795a0f724ddad7716 Mon Sep 17 00:00:00 2001
From: dewmini <5234594+dewmini@users.noreply.github.com>
Date: Mon, 20 Nov 2023 16:29:39 +1100
Subject: [PATCH] Fix issues

---
 .../grails-app/conf/application.yml           |  2 +-
 .../userdetails/cognito/Application.groovy    |  4 +++-
 .../ala/userdetails/CognitoUserService.groovy | 21 +++++++++++++++----
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/userdetails-cognito/grails-app/conf/application.yml b/userdetails-cognito/grails-app/conf/application.yml
index 4e36f2f2..cb2f54ff 100644
--- a/userdetails-cognito/grails-app/conf/application.yml
+++ b/userdetails-cognito/grails-app/conf/application.yml
@@ -187,7 +187,7 @@ environments:
 attributes:
   affiliations:
     enabled: true
-    attribute-name: 'custom:affiliation'
+    attribute-name: 'affiliation'
 
 cognito:
   mapping:
diff --git a/userdetails-cognito/grails-app/init/au/org/ala/userdetails/cognito/Application.groovy b/userdetails-cognito/grails-app/init/au/org/ala/userdetails/cognito/Application.groovy
index 6686598f..e13de238 100644
--- a/userdetails-cognito/grails-app/init/au/org/ala/userdetails/cognito/Application.groovy
+++ b/userdetails-cognito/grails-app/init/au/org/ala/userdetails/cognito/Application.groovy
@@ -83,7 +83,8 @@ class Application extends GrailsAutoConfiguration {
     }
 
     @Bean('userService')
-    IUserService userService(TokenService tokenService, EmailService emailService, AWSCognitoIdentityProvider cognitoIdp, JwtProperties jwtProperties) {
+    IUserService userService(TokenService tokenService, EmailService emailService, AWSCognitoIdentityProvider cognitoIdp, JwtProperties jwtProperties,
+                             LocationService locationService) {
 
         CognitoUserService userService = new CognitoUserService()
         userService.cognitoIdp = cognitoIdp
@@ -92,6 +93,7 @@ class Application extends GrailsAutoConfiguration {
         userService.emailService = emailService
         userService.tokenService = tokenService
         userService.jwtProperties = jwtProperties
+        userService.locationService = locationService
 
         userService.affiliationsEnabled = grailsApplication.config.getProperty('attributes.affiliations.enabled', Boolean, false)
         userService.socialLoginGroups = grailsApplication.config.getProperty('users.delegated-group-names', List, [])
diff --git a/userdetails-cognito/src/main/groovy/au/org/ala/userdetails/CognitoUserService.groovy b/userdetails-cognito/src/main/groovy/au/org/ala/userdetails/CognitoUserService.groovy
index f4c7e692..e4c99226 100644
--- a/userdetails-cognito/src/main/groovy/au/org/ala/userdetails/CognitoUserService.groovy
+++ b/userdetails-cognito/src/main/groovy/au/org/ala/userdetails/CognitoUserService.groovy
@@ -472,8 +472,8 @@ class CognitoUserService implements IUserService<UserRecord, UserPropertyRecord,
             token = results.getPaginationToken()
 
             users.each {
-                def value = it.attributes.find { it.name == "custom.$s" }?.value
-                counts[value ?: ''] = ((counts[value ?: '']) ?: 0)++
+                def value = it.attributes.find {att ->  att.name == "custom:$s" }?.value
+                counts[value ?: ''] = ((counts[value ?: '']) ?: 0) +1
             }
 
             results = token ? cognitoIdp.listUsers(new ListUsersRequest().withUserPoolId(poolId).withPaginationToken(token)) : null
@@ -733,8 +733,21 @@ class CognitoUserService implements IUserService<UserRecord, UserPropertyRecord,
             propList.addAll(userRecord.userProperties.findAll { it.name == attribute })
         }
         else if(attribute){
-            //cannot implement this since cognito does not support custom attribute search
-            throw new NotImplementedException()
+            def token
+            def results = cognitoIdp.listUsers(new ListUsersRequest().withUserPoolId(poolId))
+
+            while (results) {
+                def users = results.getUsers()
+                token = results.getPaginationToken()
+
+                users.each {
+                    def value = it.attributes.find {att ->  att.name == "custom:$attribute" }?.value
+                    if(value) {
+                        propList.add(new UserPropertyRecord(user: cognitoUserTypeToUserRecord(it, false), name: attribute, value: value))
+                    }
+                }
+                results = token ? cognitoIdp.listUsers(new ListUsersRequest().withUserPoolId(poolId).withPaginationToken(token)) : null
+            }
         }
         else{
             //cannot implement this since cognito does not support custom attribute search