Support downloading password-protected transfers

transferwee.py

@@ -41,6 +41,7 @@
 from typing import Any, Dict, List, Optional, Union
 import binascii
 import functools
+import getpass
 import hashlib
 import json
 import logging
@@ -66,7 +67,7 @@
 logger = logging.getLogger(__name__)
 
 
-def download_url(url: str) -> Optional[str]:
+def download_url(url: str, password: Optional[str] = None) -> Optional[str]:
     """Given a wetransfer.com download URL download return the downloadable URL.
 
     The URL should be of the form `https://we.tl/' or
@@ -85,6 +86,9 @@ def download_url(url: str) -> Optional[str]:
         received via email by recipients when the files are shared via email
         upload
 
+    If transfer is password-protected, the password should be provided via the
+    optional `password` argument.
+
     Return the download URL (AKA `direct_link') as a str or None if the URL
     could not be parsed.
     """
@@ -114,6 +118,11 @@ def download_url(url: str) -> Optional[str]:
         "intent": "entire_transfer",
         "security_hash": security_hash,
     }
+    if password:
+        j = {
+            **j,
+            "password": password,
+        }
     if recipient_id:
         j["recipient_id"] = recipient_id
     s = _prepare_session()
@@ -123,6 +132,14 @@ def download_url(url: str) -> Optional[str]:
     _close_session(s)
 
     j = r.json()
+    if "message" in j and j["message"] == "invalid_transfer_password":
+        if password:
+            logger.error("The provided password is incorrect")
+        else:
+            logger.error(
+                "The file is password-protected but the password hasn't been provided"
+            )
+        return None
     return j.get("direct_link")
 
 
@@ -140,16 +157,19 @@ def _file_unquote(file: str) -> str:
     )
 
 
-def download(url: str, file: str = "") -> None:
+def download(url: str, password: Optional[str] = None, file: str = "") -> None:
     """Given a `we.tl/' or `wetransfer.com/downloads/' download it.
 
+    If transfer is password-protected, the password should be provided via the
+    optional `password` argument.
+
     First a direct link is retrieved (via download_url()), the filename can be
     provided via the optional `file' argument. If not provided the filename
     will be extracted to it and it will be fetched and stored on the current
     working directory.
     """
     logger.debug(f"Downloading {url}")
-    dl_url = download_url(url)
+    dl_url = download_url(url, password)
     if not dl_url:
         logger.error(f"Could not find direct link of {url}")
         return None
@@ -620,6 +640,11 @@ def upload(
         action="store_true",
         help="only print the direct link (without downloading it)",
     )
+    dp.add_argument(
+        "-p",
+        action="store_true",
+        help="If the WETRANSFER_PASSORD env variable exists, read the transfer password from it. Otherwise, ask the user for the password.",
+    )
     dp.add_argument(
         "-o",
         type=str,
@@ -671,12 +696,19 @@ def upload(
         log.setLevel(logging.DEBUG)
 
     if args.action == "download":
+        password: Optional[str] = None
+        if args.p:
+            password = (
+                os.getenv("WETRANSFER_PASSWORD", default=None)
+                or getpass.getpass()
+            )
+
         if args.g:
             for u in args.url:
-                print(download_url(u))
+                print(download_url(u, password))
         else:
             for u in args.url:
-                download(u, args.o)
+                download(u, password, args.o)
         exit(0)
 
     if args.action == "upload":