update to pyicloud-1.0.0 to support 2fa verification

README.md

@@ -46,6 +46,7 @@ Options:
   -u, --username <username>       Your iCloud username or email address
   -p, --password <password>       Your iCloud password (default: use PyiCloud
                                   keyring or prompt for password)
+  --china-mainland                The country/region setting of your Apple ID is China mainland(中国大陆)
   --cookie-directory </cookie/directory>
                                   Directory to store cookies for
                                   authentication (default: ~/.pyicloud)

icloud.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-from pyicloud_ipd.cmdline import main
+from pyicloud.cmdline import main
 
 if __name__ == "__main__":
     main()

icloudpd/authentication.py

@@ -2,7 +2,8 @@
 
 import sys
 import click
-import pyicloud_ipd
+import pyicloud
+from pyicloud.exceptions import PyiCloudNoStoredPasswordAvailableException
 from icloudpd.logger import setup_logger
 
 
@@ -18,36 +19,60 @@ def authenticate(
         password,
         cookie_directory=None,
         raise_error_on_2sa=False,
-        client_id=None
+        client_id=None,
+        china_mainland=False,
 ):
     """Authenticate with iCloud username and password"""
     logger = setup_logger()
     logger.debug("Authenticating...")
     try:
         # If password not provided on command line variable will be set to None
-        # and PyiCloud will attempt to retrieve from its keyring
-        icloud = pyicloud_ipd.PyiCloudService(
+        # and PyiCloud will attempt to retrieve from it's keyring
+        icloud = pyicloud.PyiCloudService(
             username, password,
             cookie_directory=cookie_directory,
-            client_id=client_id)
-    except pyicloud_ipd.exceptions.NoStoredPasswordAvailable:
+            client_id=client_id,
+            china_mainland=china_mainland,
+        )
+    except PyiCloudNoStoredPasswordAvailableException:
         # Prompt for password if not stored in PyiCloud's keyring
         password = click.prompt("iCloud Password", hide_input=True)
-        icloud = pyicloud_ipd.PyiCloudService(
+        icloud = pyicloud.PyiCloudService(
             username, password,
             cookie_directory=cookie_directory,
             client_id=client_id)
-
-    if icloud.requires_2sa:
+    if icloud.requires_2fa:
+        if raise_error_on_2sa:
+            raise TwoStepAuthRequiredError(
+                "Two-factor authentication is required!"
+            )
+        logger.info("Two-factor authentication is required!")
+        request_2fa(icloud, logger)
+    elif icloud.requires_2sa:
         if raise_error_on_2sa:
             raise TwoStepAuthRequiredError(
-                "Two-step/two-factor authentication is required!"
+                "Two-step authentication is required!"
             )
-        logger.info("Two-step/two-factor authentication is required!")
+        logger.info("Two-step authentication is required!")
         request_2sa(icloud, logger)
     return icloud
 
 
+def request_2fa(icloud, logger):
+    """Request two-factor authentication."""
+    code = click.prompt("Please enter two-factor authentication code")
+    if not icloud.validate_2fa_code(code):
+        logger.error("Failed to verify two-factor authentication code")
+        sys.exit(1)
+    logger.info(
+        "Great, you're all set up. The script can now be run without "
+        "user interaction until 2FA expires.\n"
+        "You can set up email notifications for when "
+        "the two-factor authentication expires.\n"
+        "(Use --help to view information about SMTP options.)"
+    )
+
+
 def request_2sa(icloud, logger):
     """Request two-step authentication. Prompts for SMS or device"""
     devices = icloud.trusted_devices

icloudpd/base.py

@@ -14,7 +14,7 @@
 from tqdm import tqdm
 from tzlocal import get_localzone
 
-from pyicloud_ipd.exceptions import PyiCloudAPIResponseError
+from pyicloud.exceptions import PyiCloudAPIResponseException
 
 from icloudpd.logger import setup_logger
 from icloudpd.authentication import authenticate, TwoStepAuthRequiredError
@@ -50,6 +50,10 @@
     "(default: use PyiCloud keyring or prompt for password)",
     metavar="<password>",
 )
+@click.option("--china-mainland",
+              help="The country/region setting of your Apple ID is China mainland(中国大陆)",
+              is_flag=True,
+)
 @click.option(
     "--cookie-directory",
     help="Directory to store cookies for authentication "
@@ -200,6 +204,7 @@ def main(
         directory,
         username,
         password,
+        china_mainland,
         cookie_directory,
         size,
         live_photo_size,
@@ -258,6 +263,7 @@ def main(
             cookie_directory,
             raise_error_on_2sa,
             client_id=os.environ.get("CLIENT_ID"),
+            china_mainland=china_mainland,
         )
     except TwoStepAuthRequiredError:
         if notification_script is not None:
@@ -279,7 +285,7 @@ def main(
     # case exit.
     try:
         photos = icloud.photos.albums[album]
-    except PyiCloudAPIResponseError as err:
+    except PyiCloudAPIResponseException as err:
         # For later: come up with a nicer message to the user. For now take the
         # exception text
         print(err)

icloudpd/download.py

@@ -6,7 +6,7 @@
 import logging
 from tzlocal import get_localzone
 from requests.exceptions import ConnectionError  # pylint: disable=redefined-builtin
-from pyicloud_ipd.exceptions import PyiCloudAPIResponseError
+from pyicloud.exceptions import PyiCloudAPIResponseException
 from icloudpd.logger import setup_logger
 
 # Import the constants object so that we can mock WAIT_SECONDS in tests
@@ -66,7 +66,7 @@ def download_media(icloud, photo, download_path, size):
             )
             break
 
-        except (ConnectionError, socket.timeout, PyiCloudAPIResponseError) as ex:
+        except (ConnectionError, socket.timeout, PyiCloudAPIResponseException) as ex:
             if "Invalid global session" in str(ex):
                 logger.tqdm_write(
                     "Session error, re-authenticating...",

requirements.txt

@@ -1,6 +1,6 @@
-pyicloud_ipd==0.10.1
+pyicloud==1.0.0
 schema==0.7.5
-click==6.7
+click==7.1.2
 python_dateutil==2.8.2
 requests==2.28.2
 tqdm==4.64.1

tests/test_authentication.py

@@ -3,7 +3,7 @@
 from vcr import VCR
 import pytest
 from click.testing import CliRunner
-import pyicloud_ipd
+import pyicloud
 from icloudpd.base import main
 from icloudpd.authentication import authenticate, TwoStepAuthRequiredError
 import inspect
@@ -19,7 +19,7 @@ def inject_fixtures(self, caplog):
     def test_failed_auth(self):
         with vcr.use_cassette("tests/vcr_cassettes/failed_auth.yml"):
             with self.assertRaises(
-                pyicloud_ipd.exceptions.PyiCloudFailedLoginException
+                pyicloud.exceptions.PyiCloudFailedLoginException
             ) as context:
                 authenticate(
                     "bad_username",

tests/test_download_live_photos.py

@@ -10,9 +10,9 @@
 from click.testing import CliRunner
 import piexif
 from piexif._exceptions import InvalidImageDataError
-from pyicloud_ipd.services.photos import PhotoAsset, PhotoAlbum, PhotosService
-from pyicloud_ipd.base import PyiCloudService
-from pyicloud_ipd.exceptions import PyiCloudAPIResponseError
+from pyicloud.services.photos import PhotoAsset, PhotoAlbum, PhotosService
+from pyicloud.base import PyiCloudService
+from pyicloud.exceptions import PyiCloudAPIResponseException
 from requests.exceptions import ConnectionError
 from icloudpd.base import main
 from tests.helpers.print_result_exception import print_result_exception

tests/test_download_photos.py

@@ -10,9 +10,9 @@
 from click.testing import CliRunner
 import piexif
 from piexif._exceptions import InvalidImageDataError
-from pyicloud_ipd.services.photos import PhotoAsset, PhotoAlbum, PhotosService
-from pyicloud_ipd.base import PyiCloudService
-from pyicloud_ipd.exceptions import PyiCloudAPIResponseError
+from pyicloud.services.photos import PhotoAsset, PhotoAlbum, PhotosService
+from pyicloud.base import PyiCloudService
+from pyicloud.exceptions import PyiCloudAPIResponseException
 from requests.exceptions import ConnectionError
 from icloudpd.base import main
 from tests.helpers.print_result_exception import print_result_exception
@@ -454,7 +454,7 @@ def test_handle_session_error_during_download(self):
             # Pass fixed client ID via environment variable
 
             def mock_raise_response_error(arg):
-                raise PyiCloudAPIResponseError("Invalid global session", 100)
+                raise PyiCloudAPIResponseException("Invalid global session", 100)
 
             with mock.patch("time.sleep") as sleep_mock:
                 with mock.patch.object(PhotoAsset, "download") as pa_download:
@@ -522,7 +522,7 @@ def test_handle_session_error_during_photo_iteration(self):
             # Pass fixed client ID via environment variable
 
             def mock_raise_response_error(offset):
-                raise PyiCloudAPIResponseError("Invalid global session", 100)
+                raise PyiCloudAPIResponseException("Invalid global session", 100)
 
             with mock.patch("time.sleep") as sleep_mock:
                 with mock.patch.object(PhotoAlbum, "photos_request") as pa_photos_request:
@@ -655,7 +655,7 @@ def test_handle_albums_error(self):
             # Pass fixed client ID via environment variable
 
             def mock_raise_response_error():
-                raise PyiCloudAPIResponseError("Api Error", 100)
+                raise PyiCloudAPIResponseException("Api Error", 100)
 
             with mock.patch.object(PhotosService, "_fetch_folders") as pa_photos_request:
                 pa_photos_request.side_effect = mock_raise_response_error

tests/test_two_step_auth.py

@@ -6,7 +6,7 @@
 import click
 from click.testing import CliRunner
 from icloudpd.base import main
-from pyicloud_ipd import PyiCloudService
+from pyicloud import PyiCloudService
 import inspect
 import shutil