diff --git a/tests/test_user_management.py b/tests/test_user_management.py index db275247..640146b2 100644 --- a/tests/test_user_management.py +++ b/tests/test_user_management.py @@ -123,6 +123,32 @@ def mock_magic_auth_challenge_response(self): return { "id": "auth_challenge_01E4ZCR3C56J083X43JQXF3JK5", } + + @pytest.fixture + def mock_enroll_auth_factor_response(self): + + return { + "authentication_challenge": { + "object": "authentication_challenge", + "id": "auth_challenge_01FVYZWQTZQ5VB6BC5MPG2EYC5", + "created_at": "2022-02-15T15:26:53.274Z", + "updated_at": "2022-02-15T15:26:53.274Z", + "expires_at": "2022-02-15T15:36:53.279Z", + "authentication_factor_id": "auth_factor_01FVYZ5QM8N98T9ME5BCB2BBMJ", + }, + "authentication_factor": { + "object": "authentication_factor", + "id": "auth_factor_01FVYZ5QM8N98T9ME5BCB2BBMJ", + "created_at": "2022-02-15T15:14:19.392Z", + "updated_at": "2022-02-15T15:14:19.392Z", + "type": "totp", + "totp": { + "qr_code": "data:image/png;base64,{base64EncodedPng}", + "secret": "NAGCCFS3EYRB422HNAKAKY3XDUORMSRF", + "uri": "otpauth://totp/FooCorp:alan.turing@foo-corp.com?secret=NAGCCFS3EYRB422HNAKAKY3XDUORMSRF&issuer=FooCorp", + } + } + } def test_create_user(self, mock_user, mock_request_method): mock_request_method("post", mock_user, 201) @@ -390,3 +416,17 @@ def test_send_magic_auth_code(self, capture_and_mock_request, mock_user): assert url[0].endswith("user_management/magic_auth/send") assert request["json"]["email"] == email assert response["id"] == "user_01H7ZGXFP5C6BBQY6Z7277ZCT0" + + def test_enroll_auth_factor(self, mock_enroll_auth_factor_response, mock_request_method): + user = "user_01H7ZGXFP5C6BBQY6Z7277ZCT0" + type = "totp" + totp_issuer="WorkOS" + email = "marcelina@foo-corp.com" + + mock_request_method("post", mock_enroll_auth_factor_response, 200) + + enroll_auth_factor = self.user_management.enroll_auth_factor( + user, type, totp_issuer, email, + ) + + assert enroll_auth_factor == mock_enroll_auth_factor_response diff --git a/workos/resources/mfa.py b/workos/resources/mfa.py index 241df3f8..4a1528e7 100644 --- a/workos/resources/mfa.py +++ b/workos/resources/mfa.py @@ -115,3 +115,42 @@ def to_dict(self): verification_response_dict = super(WorkOSChallengeVerification, self).to_dict() return verification_response_dict + +class WorkOSAuthenticationChallengeAndFactor(WorkOSBaseResource): + """Representation of an Authentication Challenge and Factor as returned by WorkOS through the User Management feature. + + Attributes: + OBJECT_FIELDS (list): List of fields a WorkOSAuthenticationChallengeAndFactor is comprised of. + """ + + OBJECT_FIELDS = [ + "authentication_challenge", + "authentication_factor", + ] + + @classmethod + def construct_from_response(cls, response): + authentication_challenge_and_factor = super(WorkOSAuthenticationChallengeAndFactor, cls).construct_from_response( + response + ) + + authentication_challenge_and_factor.authentication_challenge = WorkOSChallenge.construct_from_response( + response["authentication_challenge"] + ) + + authentication_challenge_and_factor.authentication_factor = WorkOSAuthenticationFactorTotp.construct_from_response( + response["authentication_factor"] + ) + + return authentication_challenge_and_factor + + def to_dict(self): + authentication_challenge_and_factor_dict = super(WorkOSAuthenticationChallengeAndFactor, self).to_dict() + + challenge_dict = self.authentication_challenge.to_dict() + authentication_challenge_and_factor_dict["authentication_challenge"] = challenge_dict + + factor_dict = self.authentication_factor.to_dict() + authentication_challenge_and_factor_dict["authentication_factor"] = factor_dict + + return authentication_challenge_and_factor_dict diff --git a/workos/user_management.py b/workos/user_management.py index 3ed27a18..9aaf01e1 100644 --- a/workos/user_management.py +++ b/workos/user_management.py @@ -2,9 +2,8 @@ from workos.resources.authentication_response import WorkOSAuthenticationResponse from workos.resources.password_challenge_response import WorkOSPasswordChallengeResponse from workos.resources.list import WorkOSListResource -from workos.resources.users import ( - WorkOSUser, -) +from workos.resources.mfa import WorkOSAuthenticationChallengeAndFactor +from workos.resources.users import WorkOSUser from workos.utils.pagination_order import Order from workos.utils.request import ( RequestHelper, @@ -25,6 +24,7 @@ USER_SEND_VERIFICATION_EMAIL_PATH = "users/{0}/send_verification_email" USER_VERIFY_EMAIL_CODE_PATH = "users/verify_email_code" USER_SEND_MAGIC_AUTH_PATH = "user_management/magic_auth/send" +USER_AUTH_FACTORS_PATH = "user_management/users/{0}/auth_factors" RESPONSE_LIMIT = 10 @@ -540,3 +540,44 @@ def send_magic_auth_code( ) return WorkOSUser.construct_from_response(response).to_dict() + + def enroll_auth_factor( + self, + user, + type, + totp_issuer=None, + totp_user=None, + ): + """Enrolls a user in a new auth factor. + + Kwargs: + user (str): The unique ID of the User to be enrolled in the auth factor. + type (str): The type of factor to enroll (Only option available is 'totp'). + totp_issuer (str): Name of the Organization (Optional) + totp_user (str): Email of user (Optional) + + Returns: + dict: AuthenticationChallengeAndFactor response from WorkOS. + """ + + if type not in ["totp"]: + raise ValueError("Type parameter must be 'totp'") + + headers = {} + + payload = { + "user_id": user, + "type": type, + "totp_issuer": totp_issuer, + "totp_user": totp_user, + } + + response = self.request_helper.request( + USER_AUTH_FACTORS_PATH, + method=REQUEST_METHOD_POST, + headers=headers, + params=payload, + token=workos.api_key, + ) + + return WorkOSAuthenticationChallengeAndFactor.construct_from_response(response).to_dict()