Cognito SRP using AWS Java SDK v2.x

Question

I need to use AWS cognito to take a token and then call a secret API. I have this python script that works. How can I convert it to Java SDK v2?

I found a sample that uses the Java SDK V1, but it was not working with v2.

from warrant.aws_srp import AWSSRP
username = "user"
password = "pass"
client_id = "xxxxxxxxxxxxxxxxx"
user_pool_id = "us-east-1_xxx123445"
region = "us-east-1"
aws = AWSSRP(username=username, password=password, pool_id=user_pool_id, client_id=client_id)
tokens = aws.authenticate_user()
print(tokens)

Answer 1

Converting the AuthenticationHelper sample class from the v1 sdk to v2 isn't too difficult. The v1 calls just need to be updated to the v2 versions. All of the SRP generation remains the same. Below is how I converted the main methods from that class to use the v2 sdk.

PerformSRPAuthentication

public String PerformSRPAuthentication(String username, String password) {
    String authresult = null;

    InitiateAuthRequest authReq = initiateUserSrpAuthRequest(username);

    try {
        AnonymousCredentialsProvider creds = AnonymousCredentialsProvider.create();
        CognitoIdentityProviderClient cognitoclient = CognitoIdentityProviderClient.builder()
            .region(Region.of(this.region))
            .credentialsProvider(creds)
            .build();
        
        InitiateAuthResponse authRes = cognitoclient.initiateAuth(authReq);
        if(authRes.challengeName().equals(ChallengeNameType.PASSWORD_VERIFIER)) {
            RespondToAuthChallengeRequest challengeRequest = userSrpAuthRequest(authRes, password);
            RespondToAuthChallengeResponse result = cognitoclient.respondToAuthChallenge(challengeRequest);
            authresult = result.authenticationResult().idToken();
        }
    } catch(Exception e) {
        System.out.println("Exception: " + e);
    }
    return authresult;
}

initiateUserSrpAuthRequest

private InitiateAuthRequest initiateUserSrpAuthRequest(String username) {

    HashMap<String, String> authParams = new HashMap<String, String>();
    authParams.put("USERNAME", username);
    authParams.put("SRP_A", this.getA().toString(16));

    InitiateAuthRequest authReq = InitiateAuthRequest.builder()
        .authFlow(AuthFlowType.USER_SRP_AUTH)
        .clientId(this.clientId).authParameters(authParams).build();
    return authReq;
}

userSrpAuthRequest

private RespondToAuthChallengeRequest userSrpAuthRequest(InitiateAuthResponse challenge,
                                                         String password
) {
    String userIdForSRP = challenge.challengeParameters().get("USER_ID_FOR_SRP");
    String usernameInternal = challenge.challengeParameters().get("USERNAME");

    BigInteger B = new BigInteger(challenge.challengeParameters().get("SRP_B"), 16);
    if (B.mod(AWSAuthenticationHelper.N).equals(BigInteger.ZERO)) {
        throw new SecurityException("SRP error, B cannot be zero");
    }

    BigInteger salt = new BigInteger(challenge.challengeParameters().get("SALT"), 16);
    byte[] key = getPasswordAuthenticationKey(userIdForSRP, password, B, salt);

    Date timestamp = new Date();
    byte[] hmac = null;
    try {
        Mac mac = Mac.getInstance("HmacSHA256");
        SecretKeySpec keySpec = new SecretKeySpec(key, "HmacSHA256");
        mac.init(keySpec);
        mac.update(this.userPoolID.split("_", 2)[1].getBytes(Charset.forName("UTF-8")));
        mac.update(userIdForSRP.getBytes(Charset.forName("UTF-8")));
        byte[] secretBlock = Base64.decode(challenge.challengeParameters().get("SECRET_BLOCK"));
        mac.update(secretBlock);
        SimpleDateFormat simpleDateFormat = new SimpleDateFormat("EEE MMM d HH:mm:ss z yyyy", Locale.US);
        simpleDateFormat.setTimeZone(new SimpleTimeZone(SimpleTimeZone.UTC_TIME, "UTC"));
        String dateString = simpleDateFormat.format(timestamp);
        byte[] dateBytes = dateString.getBytes(Charset.forName("UTF-8"));
        hmac = mac.doFinal(dateBytes);
    } catch (Exception e) {
        System.out.println(e);
    }

    SimpleDateFormat formatTimestamp = new SimpleDateFormat("EEE MMM d HH:mm:ss z yyyy", Locale.US);
    formatTimestamp.setTimeZone(new SimpleTimeZone(SimpleTimeZone.UTC_TIME, "UTC"));

    Map<String, String> srpAuthResponses = new HashMap<>();
    srpAuthResponses.put("PASSWORD_CLAIM_SECRET_BLOCK", challenge.challengeParameters().get("SECRET_BLOCK"));
    srpAuthResponses.put("PASSWORD_CLAIM_SIGNATURE", new String(Base64.encode(hmac), Charset.forName("UTF-8")));
    srpAuthResponses.put("TIMESTAMP", formatTimestamp.format(timestamp));
    srpAuthResponses.put("USERNAME", usernameInternal);

    RespondToAuthChallengeRequest authChallengeRequest = RespondToAuthChallengeRequest.builder()
        .challengeName(challenge.challengeName())
        .clientId(clientId)
        .session(challenge.session())
        .challengeResponses(srpAuthResponses).build();

    return authChallengeRequest;
}

You'll also need to replace the two v1 utility classes, com.amazonaws.util.Base64 and com.amazonaws.util.StringUtils. All of the instances of StringUtils.UTF8 can be replaced with java.nio.charset.Charset.forName("UTF-8"). Replacing Base64 is a bit trickier. I ended up just copying the 4 related classes into my project locally. Those being

com.amazonaws.util.Base64 
com.amazonaws.util.CodecUtils
com.amazonaws.util.Codec
com.amazonaws.util.Base64Codec

Its not pretty, but it works for me. I'm not sure why AWS can't implement a wrapper for the SRP auth like all the other SDK's.