feat(Auth): Adding TOTP state machine actions (#3044)

* feat(Auth): Adding TOTP tasks and requests to AWSAuthCognitoPlugin

* feat(Auth): Adding TOTP state machine actions

* working on review comment

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/SignIn/SoftwareTokenSetup/CompleteTOTPSetup.swift

@@ -0,0 +1,141 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Amplify
+import AWSCognitoIdentityProvider
+
+struct CompleteTOTPSetup: Action {
+
+    var identifier: String = "CompleteTOTPSetup"
+    let userSession: String
+    let signInEventData: SignInEventData
+
+    func execute(withDispatcher dispatcher: EventDispatcher, environment: Environment) async {
+        logVerbose("\(#fileID) Starting execution", environment: environment)
+
+        do {
+            var deviceMetadata = DeviceMetadata.noData
+            guard let username = signInEventData.username else {
+                throw SignInError.unknown(message: "Unable to unwrap username during TOTP verification")
+            }
+            let authEnv = try environment.authEnvironment()
+            let userpoolEnv = try environment.userPoolEnvironment()
+            let challengeType: CognitoIdentityProviderClientTypes.ChallengeNameType = .mfaSetup
+
+            deviceMetadata = await DeviceMetadataHelper.getDeviceMetadata(
+                for: username,
+                with: environment)
+
+            var challengeResponses = [
+                "USERNAME": username
+            ]
+            let userPoolClientId = userpoolEnv.userPoolConfiguration.clientId
+
+            if let clientSecret = userpoolEnv.userPoolConfiguration.clientSecret {
+                let clientSecretHash = ClientSecretHelper.clientSecretHash(
+                    username: username,
+                    userPoolClientId: userPoolClientId,
+                    clientSecret: clientSecret
+                )
+                challengeResponses["SECRET_HASH"] = clientSecretHash
+            }
+
+            if case .metadata(let data) = deviceMetadata {
+                challengeResponses["DEVICE_KEY"] = data.deviceKey
+            }
+
+            let asfDeviceId = try await CognitoUserPoolASF.asfDeviceID(
+                for: username,
+                credentialStoreClient: authEnv.credentialsClient)
+
+            var userContextData: CognitoIdentityProviderClientTypes.UserContextDataType?
+            if let encodedData = CognitoUserPoolASF.encodedContext(
+                username: username,
+                asfDeviceId: asfDeviceId,
+                asfClient: userpoolEnv.cognitoUserPoolASFFactory(),
+                userPoolConfiguration: userpoolEnv.userPoolConfiguration) {
+                userContextData = .init(encodedData: encodedData)
+            }
+
+            let analyticsMetadata = userpoolEnv
+                .cognitoUserPoolAnalyticsHandlerFactory()
+                .analyticsMetadata()
+
+            let input = RespondToAuthChallengeInput(
+                analyticsMetadata: analyticsMetadata,
+                challengeName: challengeType,
+                challengeResponses: challengeResponses,
+                clientId: userPoolClientId,
+                session: userSession,
+                userContextData: userContextData)
+
+            let responseEvent = try await UserPoolSignInHelper.sendRespondToAuth(
+                request: input,
+                for: username,
+                signInMethod: signInEventData.signInMethod,
+                environment: userpoolEnv)
+            logVerbose("\(#fileID) Sending event \(responseEvent)",
+                       environment: environment)
+            await dispatcher.send(responseEvent)
+            // TODO: HS:
+            //        } catch let error where deviceNotFound(error: error, deviceMetadata: deviceMetadata) {
+            //            logVerbose("\(#fileID) Received device not found \(error)", environment: environment)
+            //            // Remove the saved device details and retry verify challenge
+            //            await DeviceMetadataHelper.removeDeviceMetaData(for: username, with: environment)
+            //            let event = SignInChallengeEvent(
+            //                eventType: .retryVerifyChallengeAnswer(confirmSignEventData)
+            //            )
+            //            logVerbose("\(#fileID) Sending event \(event)", environment: environment)
+            //            await dispatcher.send(event)
+        } catch let error as SignInError {
+            logError(error.authError.errorDescription, environment: environment)
+            let errorEvent = SignInEvent(eventType: .throwAuthError(error))
+            logVerbose("\(#fileID) Sending event \(errorEvent)",
+                       environment: environment)
+            await dispatcher.send(errorEvent)
+        } catch {
+            let error = SignInError.service(error: error)
+            logError(error.authError.errorDescription, environment: environment)
+            let errorEvent = SignInEvent(eventType: .throwAuthError(error))
+            logVerbose("\(#fileID) Sending event \(errorEvent)",
+                       environment: environment)
+            await dispatcher.send(errorEvent)
+        }
+    }
+
+// TODO: HS: Figure out if this is needed
+//    func deviceNotFound(error: Error, deviceMetadata: DeviceMetadata) -> Bool {
+//
+//        // If deviceMetadata was not send, the error returned is not from device not found.
+//        if case .noData = deviceMetadata {
+//            return false
+//        }
+//
+//        if let serviceError: RespondToAuthChallengeOutputError = error.internalAWSServiceError(),
+//           case .resourceNotFoundException = serviceError {
+//            return true
+//        }
+//        return false
+//    }
+
+}
+
+extension CompleteTOTPSetup: CustomDebugDictionaryConvertible {
+    var debugDictionary: [String: Any] {
+        [
+            "identifier": identifier,
+            "session": userSession.masked(),
+            "signInEventData": signInEventData.debugDictionary
+        ]
+    }
+}
+
+extension CompleteTOTPSetup: CustomDebugStringConvertible {
+    var debugDescription: String {
+        debugDictionary.debugDescription
+    }
+}

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/SignIn/SoftwareTokenSetup/InitializeTOTPSetup.swift

@@ -0,0 +1,40 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+
+struct InitializeTOTPSetup: Action {
+
+    var identifier: String = "InitializeTOTPSetup"
+    let authResponse: SignInResponseBehavior
+
+    func execute(withDispatcher dispatcher: EventDispatcher, environment: Environment) async {
+        logVerbose("\(#fileID) Start execution", environment: environment)
+        let event = SetUpTOTPEvent(
+            id: UUID().uuidString,
+            eventType: .setUpTOTP(authResponse))
+        logVerbose("\(#fileID) Sending event \(event.type)", environment: environment)
+        await dispatcher.send(event)
+    }
+}
+
+extension InitializeTOTPSetup: CustomDebugDictionaryConvertible {
+    var debugDictionary: [String: Any] {
+        [
+            "identifier": identifier,
+            "challengeName": authResponse.challengeName?.rawValue ?? "",
+            "session": authResponse.session?.masked() ?? "",
+            "challengeParameters": authResponse.challengeParameters ?? [:]
+        ]
+    }
+}
+
+extension InitializeTOTPSetup: CustomDebugStringConvertible {
+    var debugDescription: String {
+        debugDictionary.debugDescription
+    }
+}

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/SignIn/SoftwareTokenSetup/SetUpTOTP.swift

@@ -0,0 +1,80 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Amplify
+import Foundation
+import AWSCognitoIdentityProvider
+
+struct SetUpTOTP: Action {
+
+    var identifier: String = "SetUpTOTP"
+    let authResponse: SignInResponseBehavior
+    let signInEventData: SignInEventData
+
+    func execute(withDispatcher dispatcher: EventDispatcher, environment: Environment) async {
+        logVerbose("\(#fileID) Starting execution", environment: environment)
+
+        do {
+            let userpoolEnv = try environment.userPoolEnvironment()
+            let client = try userpoolEnv.cognitoUserPoolFactory()
+            let input = AssociateSoftwareTokenInput(session: authResponse.session)
+
+            // Initiate Set Up TOTP
+            let result = try await client.associateSoftwareToken(input: input)
+
+            guard let username = signInEventData.username else {
+                throw SignInError.unknown(message: "Unable unwrap username to for use during TOTP setup")
+            }
+
+            guard let session = result.session,
+                  let secretCode = result.secretCode else {
+                throw SignInError.unknown(message: "Error unwrapping result associateSoftwareToken result")
+            }
+
+            let responseEvent = SetUpTOTPEvent(eventType:
+                    .waitForAnswer(.init(
+                        secretCode: secretCode,
+                        session: session,
+                        username: username)))
+            logVerbose("\(#fileID) Sending event \(responseEvent)",
+                       environment: environment)
+            await dispatcher.send(responseEvent)
+        } catch let error as SignInError {
+            logError(error.authError.errorDescription, environment: environment)
+            let errorEvent = SetUpTOTPEvent(eventType: .throwError(error))
+            logVerbose("\(#fileID) Sending event \(errorEvent)",
+                       environment: environment)
+            await dispatcher.send(errorEvent)
+        } catch {
+            let error = SignInError.service(error: error)
+            logError(error.authError.errorDescription, environment: environment)
+            let errorEvent = SetUpTOTPEvent(eventType: .throwError(error))
+            logVerbose("\(#fileID) Sending event \(errorEvent)",
+                       environment: environment)
+            await dispatcher.send(errorEvent)
+        }
+    }
+
+}
+
+extension SetUpTOTP: CustomDebugDictionaryConvertible {
+    var debugDictionary: [String: Any] {
+        [
+            "identifier": identifier,
+            "challengeName": authResponse.challengeName?.rawValue ?? "",
+            "session": authResponse.session?.masked() ?? "",
+            "challengeParameters": authResponse.challengeParameters ?? [:],
+            "signInEventData": signInEventData.debugDictionary
+        ]
+    }
+}
+
+extension SetUpTOTP: CustomDebugStringConvertible {
+    var debugDescription: String {
+        debugDictionary.debugDescription
+    }
+}

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/SignIn/SoftwareTokenSetup/VerifyTOTPSetup.swift

@@ -0,0 +1,75 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Amplify
+import Foundation
+import AWSCognitoIdentityProvider
+
+struct VerifyTOTPSetup: Action {
+
+    var identifier: String = "VerifyTOTPSetup"
+    let session: String
+    let totpCode: String
+    let friendlyDeviceName: String?
+
+    func execute(withDispatcher dispatcher: EventDispatcher, environment: Environment) async {
+        logVerbose("\(#fileID) Starting execution", environment: environment)
+
+        do {
+            let userpoolEnv = try environment.userPoolEnvironment()
+            let client = try userpoolEnv.cognitoUserPoolFactory()
+            let input = VerifySoftwareTokenInput(
+                friendlyDeviceName: friendlyDeviceName,
+                session: session,
+                userCode: totpCode)
+
+            // Initiate TOTP verification
+            let result = try await client.verifySoftwareToken(input: input)
+
+            guard let session = result.session else {
+                throw SignInError.unknown(message: "Unable to retrieve the session value from VerifySoftwareToken response")
+            }
+
+            let responseEvent = SetUpTOTPEvent(eventType:
+                    .respondToAuthChallenge(session))
+            logVerbose("\(#fileID) Sending event \(responseEvent)",
+                       environment: environment)
+            await dispatcher.send(responseEvent)
+        } catch let error as SignInError {
+            logError(error.authError.errorDescription, environment: environment)
+            let errorEvent = SetUpTOTPEvent(eventType: .throwError(error))
+            logVerbose("\(#fileID) Sending event \(errorEvent)",
+                       environment: environment)
+            await dispatcher.send(errorEvent)
+        } catch {
+            let error = SignInError.service(error: error)
+            logError(error.authError.errorDescription, environment: environment)
+            let errorEvent = SetUpTOTPEvent(eventType: .throwError(error))
+            logVerbose("\(#fileID) Sending event \(errorEvent)",
+                       environment: environment)
+            await dispatcher.send(errorEvent)
+        }
+    }
+
+}
+
+extension VerifyTOTPSetup: CustomDebugDictionaryConvertible {
+    var debugDictionary: [String: Any] {
+        [
+            "identifier": identifier,
+            "session": session.masked(),
+            "totpCode": totpCode.redacted(),
+            "friendlyDeviceName": friendlyDeviceName ?? ""
+        ]
+    }
+}
+
+extension VerifyTOTPSetup: CustomDebugStringConvertible {
+    var debugDescription: String {
+        debugDictionary.debugDescription
+    }
+}

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Request/FetchMFAPreferenceRequest.swift

@@ -13,7 +13,7 @@ public struct FetchMFAPreferenceRequest: AmplifyOperationRequest {
     /// Extra request options defined in `FetchMFAPreferenceRequest.Options`
     public var options: Options
 
-    public init(options: Options) {
+    internal init(options: Options) {
         self.options = options
     }
 }

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Request/UpdateMFAPreferenceRequest.swift

@@ -11,14 +11,13 @@ import Foundation
 /// Request for udpating user MFA preferences
 public struct UpdateMFAPreferenceRequest: AmplifyOperationRequest {
 
-    public let smsPreference: MFAPreference?
-
-    public let totpPreference: MFAPreference?
+    internal let smsPreference: MFAPreference?
+    internal let totpPreference: MFAPreference?
 
     /// Extra request options defined in `UpdateMFAPreferenceRequest.Options`
     public var options: Options
 
-    public init(options: Options,
+    internal init(options: Options,
                   smsPreference: MFAPreference?,
                   totpPreference: MFAPreference?) {
         self.options = options

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Support/Constants/AuthPluginErrorConstants.swift

@@ -214,6 +214,12 @@ extension AuthPluginErrorConstants {
         "Make sure that a valid challenge response is passed for confirmSignIn"
     )
 
+    static let confirmSignInMFASelectionResponseError: AuthPluginValidationErrorString = (
+        "challengeResponse",
+        "challengeResponse for MFA selection can only have SMS_MFA or SOFTWARE_TOKEN_MFA.",
+        "Make sure that a valid challenge response is passed for confirmSignIn. Try using `MFAType.totp.challengeResponse` or `MFAType.sms.challengeResponse` as the challenge response"
+    )
+
     static let confirmResetPasswordUsernameError: AuthPluginValidationErrorString = (
         "username",
         "username is required to confirmResetPassword",
@@ -332,4 +338,8 @@ extension AuthPluginErrorConstants {
     Check if you are allowed to make this request based on the web ACL thats associated with your user pool 
     """
 
+    static let concurrentModificationException: RecoverySuggestion = """
+    Make sure the requests sent are controlled and concurrent operations are handled properly
+    """
+
 }