From 32ac85878ea0021a75b4bbef229d5f4cd9cfde1d Mon Sep 17 00:00:00 2001 From: moeez96 Date: Fri, 17 Mar 2023 12:11:13 +0500 Subject: [PATCH] refactor: Resume-able Apple migration mgmt cmds --- .../generate_and_store_apple_transfer_ids.py | 74 ++++++++++++++----- .../generate_and_store_new_apple_ids.py | 65 ++++++++++++---- 2 files changed, 109 insertions(+), 30 deletions(-) diff --git a/common/djangoapps/third_party_auth/management/commands/generate_and_store_apple_transfer_ids.py b/common/djangoapps/third_party_auth/management/commands/generate_and_store_apple_transfer_ids.py index f7856d3cfc0b..d367b30ae2a2 100644 --- a/common/djangoapps/third_party_auth/management/commands/generate_and_store_apple_transfer_ids.py +++ b/common/djangoapps/third_party_auth/management/commands/generate_and_store_apple_transfer_ids.py @@ -10,6 +10,7 @@ from django.core.management.base import BaseCommand, CommandError from django.db import transaction +from django.db.models import Q import jwt from social_django.models import UserSocialAuth from social_django.utils import load_strategy @@ -21,6 +22,12 @@ log = logging.getLogger(__name__) +class AccessTokenExpiredException(Exception): + """ + Raised when access token has been expired. + """ + + class Command(BaseCommand): """ Management command to generate transfer identifiers for apple users using their apple_id @@ -78,37 +85,70 @@ def _generate_access_token(self, client_secret): access_token = response.json().get('access_token') return access_token - def add_arguments(self, parser): - parser.add_argument('target_team_id', help='Team ID to which the app is to be migrated to.') + def _get_token_and_secret(self): + """ + Get access_token and client_secret + """ + client_secret = self._generate_client_secret() + access_token = self._generate_access_token(client_secret) + return access_token, client_secret - @transaction.atomic - def handle(self, *args, **options): - target_team_id = options['target_team_id'] + def _update_token_and_secret(self): + self.access_token, self.client_secret = self._get_token_and_secret() # pylint: disable=W0201 + def _fetch_transfer_id(self, apple_id, target_team_id): + """ + Fetch Transfer ID for a given Apple ID from Apple API. + """ migration_url = "https://appleid.apple.com/auth/usermigrationinfo" app_id = "org.edx.mobile" - - client_secret = self._generate_client_secret() - access_token = self._generate_access_token(client_secret) - if not access_token: - raise CommandError('Failed to create access token.') - headers = { "Content-Type": "application/x-www-form-urlencoded", "Host": "appleid.apple.com", - "Authorization": "Bearer " + access_token + "Authorization": "Bearer " + self.access_token } payload = { "target": target_team_id, "client_id": app_id, - "client_secret": client_secret + "client_secret": self.client_secret, + "sub": apple_id } + response = requests.post(migration_url, data=payload, headers=headers) + if response.status_code == 400: + raise AccessTokenExpiredException + + return response.json().get('transfer_sub') + + def _get_transfer_id_for_apple_id(self, apple_id, target_team_id): + """ + Given an Apple ID from the old transferring team, + create and return its respective transfer id. + """ + try: + transfer_id = self._fetch_transfer_id(apple_id, target_team_id) + except AccessTokenExpiredException: + log.info('Access token expired. Re-creating access token.') + self._update_token_and_secret() + transfer_id = self._fetch_transfer_id(apple_id, target_team_id) + return transfer_id + + def add_arguments(self, parser): + parser.add_argument('target_team_id', help='Team ID to which the app is to be migrated to.') + + @transaction.atomic + def handle(self, *args, **options): + target_team_id = options['target_team_id'] + + self._update_token_and_secret() + if not self.access_token: + raise CommandError('Failed to create access token.') - apple_ids = UserSocialAuth.objects.filter(provider=AppleIdAuth.name).values_list('uid', flat=True) + already_processed_apple_ids = AppleMigrationUserIdInfo.objects.all().exclude( + Q(transfer_id__isnull=True) | Q(transfer_id="")).values_list('old_apple_id', flat=True) + apple_ids = UserSocialAuth.objects.filter(provider=AppleIdAuth.name).exclude( + uid__in=already_processed_apple_ids).values_list('uid', flat=True) for apple_id in apple_ids: - payload['sub'] = apple_id - response = requests.post(migration_url, data=payload, headers=headers) - transfer_id = response.json().get('transfer_sub') + transfer_id = self._get_transfer_id_for_apple_id(apple_id, target_team_id) if transfer_id: apple_user_id_info, _ = AppleMigrationUserIdInfo.objects.get_or_create(old_apple_id=apple_id) apple_user_id_info.transfer_id = transfer_id diff --git a/common/djangoapps/third_party_auth/management/commands/generate_and_store_new_apple_ids.py b/common/djangoapps/third_party_auth/management/commands/generate_and_store_new_apple_ids.py index 5f9c69f779bc..1acb475d0e94 100644 --- a/common/djangoapps/third_party_auth/management/commands/generate_and_store_new_apple_ids.py +++ b/common/djangoapps/third_party_auth/management/commands/generate_and_store_new_apple_ids.py @@ -10,6 +10,7 @@ from django.core.management.base import BaseCommand, CommandError from django.db import transaction +from django.db.models import Q import jwt from social_django.utils import load_strategy @@ -19,6 +20,12 @@ log = logging.getLogger(__name__) +class AccessTokenExpiredException(Exception): + """ + Raised when access token has been expired. + """ + + class Command(BaseCommand): """ Management command to exchange transfer identifiers for new team-scoped identifier for @@ -76,31 +83,63 @@ def _generate_access_token(self, client_secret): access_token = response.json().get('access_token') return access_token - @transaction.atomic - def handle(self, *args, **options): - migration_url = "https://appleid.apple.com/auth/usermigrationinfo" - app_id = "org.edx.mobile" - + def _get_token_and_secret(self): + """ + Get access_token and client_secret + """ client_secret = self._generate_client_secret() access_token = self._generate_access_token(client_secret) - if not access_token: - raise CommandError('Failed to create access token.') + return access_token, client_secret + + def _update_token_and_secret(self): + self.access_token, self.client_secret = self._get_token_and_secret() # pylint: disable=W0201 + def _fetch_new_apple_id(self, transfer_id): + """ + Fetch Apple ID for a given transfer ID from Apple API. + """ + migration_url = "https://appleid.apple.com/auth/usermigrationinfo" + app_id = "org.edx.mobile" headers = { "Content-Type": "application/x-www-form-urlencoded", "Host": "appleid.apple.com", - "Authorization": "Bearer " + access_token + "Authorization": "Bearer " + self.access_token } payload = { "client_id": app_id, - "client_secret": client_secret + "client_secret": self.client_secret, + "transfer_sub": transfer_id } + response = requests.post(migration_url, data=payload, headers=headers) + if response.status_code == 400: + raise AccessTokenExpiredException + + return response.json().get('sub') + + def _exchange_transfer_id_for_new_apple_id(self, transfer_id): + """ + For a Transfer ID obtained from the transferring team, + return the correlating Apple ID belonging to the recipient team. + """ + try: + new_apple_id = self._fetch_new_apple_id(transfer_id) + except AccessTokenExpiredException: + log.info('Access token expired. Re-creating access token.') + self._update_token_and_secret() + new_apple_id = self._fetch_new_apple_id(transfer_id) + + return new_apple_id + + @transaction.atomic + def handle(self, *args, **options): + self._update_token_and_secret() + if not self.access_token: + raise CommandError('Failed to create access token.') - apple_user_ids_info = AppleMigrationUserIdInfo.objects.all() + apple_user_ids_info = AppleMigrationUserIdInfo.objects.filter(Q(new_apple_id__isnull=True) | Q(new_apple_id=""), + ~Q(transfer_id=""), transfer_id__isnull=False) for apple_user_id_info in apple_user_ids_info: - payload['transfer_sub'] = apple_user_id_info.transfer_id - response = requests.post(migration_url, data=payload, headers=headers) - new_apple_id = response.json().get('sub') + new_apple_id = self._exchange_transfer_id_for_new_apple_id(apple_user_id_info.transfer_id) if new_apple_id: apple_user_id_info.new_apple_id = new_apple_id apple_user_id_info.save()