feat(Auth): Adding TOTP states, events, data models and resolvers

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/Actions/SignIn/DeviceSRPAuth/InitiateAuthDeviceSRP.swift

@@ -86,6 +86,8 @@ struct InitiateAuthDeviceSRP: Action {
         _ response: SignInResponseBehavior,
         with stateData: SRPStateData) -> StateMachineEvent {
 
+
+            //HS: TODO: Combine this logic with responedToAuthChallenge
             if let challengeName = response.challengeName {
                 switch challengeName {
                 case .devicePasswordVerifier:

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/Data/ConfirmSignInEventData.swift

@@ -12,6 +12,7 @@ struct ConfirmSignInEventData {
     let answer: String
     let attributes: [String: String]
     let metadata: [String: String]?
+    let friendlyDeviceName: String?
 
 }
 

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/Data/RespondToAuthChallenge.swift

@@ -39,18 +39,45 @@ extension RespondToAuthChallenge {
                                        attributeKey: nil)
     }
 
+    var getAllowedMFATypesForSelection: Set<MFAType> {
+        return getMFATypes(forKey: "MFAS_CAN_CHOOSE")
+    }
+
+    var getAllowedMFATypesForSetup: Set<MFAType> {
+        return getMFATypes(forKey: "MFAS_CAN_SETUP")
+    }
+
+    /// Helper method to extract MFA types from parameters
+    private func getMFATypes(forKey key: String) -> Set<MFAType> {
+        var mfaTypes = Set<MFAType>()
+        guard let mfaTypeParametersData = parameters?[key]?.data(using: .utf8),
+              let mfaTypesArray = try? JSONDecoder().decode(
+                [String].self, from: mfaTypeParametersData) else {
+            return mfaTypes
+        }
+
+        for mfaTypeValue in mfaTypesArray {
+            if let mfaType = MFAType(rawValue: String(mfaTypeValue)) {
+                mfaTypes.insert(mfaType)
+            }
+        }
+
+        return mfaTypes
+    }
+
     var debugDictionary: [String: Any] {
         return ["challenge": challenge,
                 "username": username.masked()]
     }
 
     func getChallengeKey() throws -> String {
         switch challenge {
-        case .customChallenge: return "ANSWER"
+        case .customChallenge, .selectMfaType: return "ANSWER"
         case .smsMfa: return "SMS_MFA_CODE"
+        case .softwareTokenMfa: return "SOFTWARE_TOKEN_MFA_CODE"
         case .newPasswordRequired: return "NEW_PASSWORD"
         default:
-            let message = "UnSupported challenge response \(challenge)"
+            let message = "Unsupported challenge type for response key generation \(challenge)"
             let error = SignInError.unknown(message: message)
             throw error
         }

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/Data/SignInTOTPSetupData.swift

@@ -0,0 +1,30 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Amplify
+
+struct SignInTOTPSetupData {
+
+    let secretCode: String
+    let session: String
+    let username: String
+
+}
+
+extension SignInTOTPSetupData: CustomDebugDictionaryConvertible {
+    var debugDictionary: [String: Any] {
+        [
+            "sharedSecret": secretCode.redacted(),
+            "session": session.masked(),
+            "username": username.masked()
+        ]
+    }
+}
+
+extension SignInTOTPSetupData: Codable { }
+
+extension SignInTOTPSetupData: Equatable { }

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/Events/SetUpTOTPEvent.swift

@@ -0,0 +1,71 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+
+/// Session value created by the service
+typealias UserSession = String
+
+struct SetUpTOTPEvent: StateMachineEvent {
+
+    enum EventType {
+
+        case setUpTOTP(SignInResponseBehavior)
+
+        case waitForAnswer(SignInTOTPSetupData)
+
+        case verifyChallengeAnswer(ConfirmSignInEventData)
+
+        case respondToAuthChallenge(UserSession)
+
+        case verified
+
+        case throwError(SignInError)
+
+    }
+
+    let id: String
+    let eventType: EventType
+    let time: Date?
+
+    var type: String {
+        switch eventType {
+        case .setUpTOTP: return "SetUpTOTPEvent.setUpTOTP"
+        case .verified: return "SetUpTOTPEvent.verified"
+        case .verifyChallengeAnswer: return "SetUpTOTPEvent.verifyChallengeAnswer"
+        case .waitForAnswer: return "SetUpTOTPEvent.waitForAnswer"
+        case .respondToAuthChallenge: return "SetUpTOTPEvent.respondToAuthChallenge"
+        case .throwError: return "SetUpTOTPEvent.throwError"
+        }
+    }
+
+    init(id: String = UUID().uuidString,
+         eventType: EventType,
+         time: Date? = nil) {
+        self.id = id
+        self.eventType = eventType
+        self.time = time
+    }
+}
+
+extension SetUpTOTPEvent.EventType: Equatable {
+    static func == (lhs: SetUpTOTPEvent.EventType, rhs: SetUpTOTPEvent.EventType) -> Bool {
+        switch (lhs, rhs) {
+        case (.setUpTOTP, .setUpTOTP),
+            (.verified, .verified),
+            (.verifyChallengeAnswer, .verifyChallengeAnswer),
+            (.waitForAnswer, .waitForAnswer),
+            (.respondToAuthChallenge, .respondToAuthChallenge),
+            (.throwError, .throwError):
+            return true
+        default:
+            return false
+        }
+    }
+
+
+}

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/Events/SignInEvent.swift

@@ -38,6 +38,8 @@ struct SignInEvent: StateMachineEvent {
 
         case respondDevicePasswordVerifier(SRPStateData, SignInResponseBehavior)
 
+        case initiateTOTPSetup(Username, SignInResponseBehavior)
+
         case throwPasswordVerifierError(SignInError)
 
         case finalizeSignIn(SignedInData)
@@ -76,6 +78,7 @@ struct SignInEvent: StateMachineEvent {
         case .receivedChallenge: return "SignInEvent.receivedChallenge"
         case .verifySMSChallenge: return "SignInEvent.verifySMSChallenge"
         case .retryRespondPasswordVerifier: return "SignInEvent.retryRespondPasswordVerifier"
+        case .initiateTOTPSetup: return "SignInEvent.initiateTOTPSetup"
         }
     }
 
@@ -109,7 +112,8 @@ extension SignInEvent.EventType: Equatable {
             (.throwAuthError, .throwAuthError),
             (.receivedChallenge, .receivedChallenge),
             (.verifySMSChallenge, .verifySMSChallenge),
-            (.retryRespondPasswordVerifier, .retryRespondPasswordVerifier):
+            (.retryRespondPasswordVerifier, .retryRespondPasswordVerifier),
+            (.initiateTOTPSetup, .initiateTOTPSetup):
             return true
         default: return false
         }

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/DebugInfo/SignInState+Debug.swift

@@ -40,6 +40,10 @@ extension SignInState {
             additionalMetadataDictionary = ["DeviceSRPState": deviceSRPState.debugDictionary]
         case .signedIn(let data):
             additionalMetadataDictionary = ["SignedInData": data.debugDictionary]
+        case .resolvingTOTPSetup(let signInTOTPSetupState, let signInEventData):
+            additionalMetadataDictionary = [
+                "SignInTOTPSetupState": signInTOTPSetupState.debugDictionary,
+                "SignInEventData" : signInEventData.debugDictionary]
         case .error:
             additionalMetadataDictionary = [:]
         }

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/DebugInfo/SignInTOTPSetupState+Debug.swift

@@ -0,0 +1,28 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+
+extension SignInTOTPSetupState {
+
+    var debugDictionary: [String: Any] {
+        var additionalMetadataDictionary: [String: Any] = [:]
+        switch self {
+        case .waitingForAnswer(let signInTOTPSetupData):
+            additionalMetadataDictionary = signInTOTPSetupData.debugDictionary
+        case .verifying(let signInSetupData, let confirmSignInEventData):
+            additionalMetadataDictionary = confirmSignInEventData.debugDictionary
+            additionalMetadataDictionary = additionalMetadataDictionary.merging(
+                signInSetupData.debugDictionary,
+                uniquingKeysWith: {$1})
+        case .error(let error):
+            additionalMetadataDictionary["error"] = error
+        default: additionalMetadataDictionary = [:]
+        }
+        return [type: additionalMetadataDictionary]
+    }
+}

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/SignInState.swift

@@ -14,6 +14,7 @@ enum SignInState: State {
     case signingInWithCustom(CustomSignInState, SignInEventData)
     case signingInViaMigrateAuth(MigrateSignInState, SignInEventData)
     case resolvingChallenge(SignInChallengeState, AuthChallengeType, SignInMethod)
+    case resolvingTOTPSetup(SignInTOTPSetupState, SignInEventData)
     case signingInWithHostedUI(HostedUISignInState)
     case confirmingDevice
     case resolvingDeviceSrpa(DeviceSRPState)
@@ -32,6 +33,7 @@ extension SignInState {
         case .signingInWithCustom: return "SignInState.signingInWithCustom"
         case .signingInViaMigrateAuth: return "SignInState.signingInViaMigrateAuth"
         case .resolvingChallenge: return "SignInState.resolvingChallenge"
+        case .resolvingTOTPSetup: return "SignInState.resolvingTOTPSetup"
         case .confirmingDevice: return "SignInState.confirmingDevice"
         case .resolvingDeviceSrpa: return "SignInState.resolvingDeviceSrpa"
         case .signedIn: return "SignInState.signedIn"

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/CodeGen/States/SignInTOTPSetupState.swift

@@ -0,0 +1,57 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+import Foundation
+
+enum SignInTOTPSetupState: State {
+
+    case notStarted
+
+    case setUpTOTP
+
+    case waitingForAnswer(SignInTOTPSetupData)
+
+    case verifying(SignInTOTPSetupData, ConfirmSignInEventData)
+
+    case respondingToAuthChallenge
+
+    case success
+
+    case error(SignInTOTPSetupData?, SignInError)
+}
+
+extension SignInTOTPSetupState {
+
+    var type: String {
+        switch self {
+        case .notStarted: return "SignInTOTPSetupState.notStarted"
+        case .setUpTOTP: return "SignInTOTPSetupState.setUpTOTP"
+        case .waitingForAnswer: return "SignInTOTPSetupState.waitingForAnswer"
+        case .verifying: return "SignInTOTPSetupState.verifying"
+        case .respondingToAuthChallenge: return "SignInTOTPSetupState.respondingToAuthChallenge"
+        case .success: return "SignInTOTPSetupState.success"
+        case .error: return "SignInTOTPSetupState.error"
+        }
+    }
+}
+
+extension SignInTOTPSetupState: Equatable {
+    static func == (lhs: SignInTOTPSetupState, rhs: SignInTOTPSetupState) -> Bool {
+        switch (lhs, rhs) {
+        case (.notStarted, .notStarted),
+            (.setUpTOTP, .setUpTOTP),
+            (.waitingForAnswer, .waitingForAnswer),
+            (.verifying, .verifying),
+            (.respondingToAuthChallenge, .respondingToAuthChallenge),
+            (.success, .success),
+            (.error, .error):
+            return true
+        default: return false
+        }
+    }
+
+}

AmplifyPlugins/Auth/Sources/AWSCognitoAuthPlugin/StateMachine/ErrorMapping/SignInError+Helper.swift

@@ -25,6 +25,12 @@ extension SignInError {
                case .userNotConfirmedException = internalError {
                 return true
             }
+
+            if let internalError: VerifySoftwareTokenOutputError = serviceError.internalAWSServiceError(),
+               case .userNotConfirmedException = internalError {
+                return true
+            }
+
         default: break
         }
         return false
@@ -42,6 +48,11 @@ extension SignInError {
                case .passwordResetRequiredException = internalError {
                 return true
             }
+
+            if let internalError: VerifySoftwareTokenOutputError = serviceError.internalAWSServiceError(),
+               case .passwordResetRequiredException = internalError {
+                return true
+            }
         default: break
         }
         return false