From dd9313c3e6da97cf09edd83dbab2386fbfa7e137 Mon Sep 17 00:00:00 2001
From: chenwenlong02 <chenwenlong02@meituan.com>
Date: Tue, 17 Jan 2023 23:40:08 +0800
Subject: [PATCH] update to pyicloud-1.0.0 to support 2fa verification

---
 icloudpd/authentication.py | 38 ++++++++++++++++++++++++++++++--------
 icloudpd/base.py           | 16 ++--------------
 icloudpd/download.py       |  4 ++--
 requirements.txt           |  6 +-----
 4 files changed, 35 insertions(+), 29 deletions(-)

diff --git a/icloudpd/authentication.py b/icloudpd/authentication.py
index e12fe6a6b..ef33f5006 100644
--- a/icloudpd/authentication.py
+++ b/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
 
 
@@ -26,28 +27,49 @@ def authenticate(
     try:
         # If password not provided on command line variable will be set to None
         # and PyiCloud will attempt to retrieve from it's keyring
-        icloud = pyicloud_ipd.PyiCloudService(
+        icloud = pyicloud.PyiCloudService(
             username, password,
             cookie_directory=cookie_directory,
             client_id=client_id)
-    except pyicloud_ipd.exceptions.NoStoredPasswordAvailable:
+    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
diff --git a/icloudpd/base.py b/icloudpd/base.py
index 1298c7356..6ed29ff86 100755
--- a/icloudpd/base.py
+++ b/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
@@ -279,7 +279,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)
@@ -364,18 +364,6 @@ def photos_exception_handler(ex, retries):
         logger.set_tqdm(photos_enumerator)
 
     def download_photo(counter, photo):
-        """internal function for actually downloading the photos"""
-        if skip_videos and photo.item_type != "image":
-            logger.set_tqdm_description(
-                "Skipping %s, only downloading photos." % photo.filename
-            )
-            return
-        if photo.item_type != "image" and photo.item_type != "movie":
-            logger.set_tqdm_description(
-                "Skipping %s, only downloading photos and videos. "
-                "(Item type was: %s)" % (photo.filename, photo.item_type)
-            )
-            return
         try:
             created_date = photo.created.astimezone(get_localzone())
         except (ValueError, OSError):
diff --git a/icloudpd/download.py b/icloudpd/download.py
index fd092329d..2bfab3cfc 100644
--- a/icloudpd/download.py
+++ b/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
@@ -65,7 +65,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...",
diff --git a/requirements.txt b/requirements.txt
index b5b6bd4d9..a59cd1488 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,7 +1,3 @@
-pyicloud_ipd==0.10.1
-schema==0.7.4
-click==6.7
-python_dateutil==2.8.2
+pyicloud==1.0.0
 requests>=2.20.0
-tqdm==4.62.0
 piexif==1.1.3