Allow specifying a custom endpoint URL for NUS downloads

src/libWiiPy/title/nus.py

@@ -6,6 +6,7 @@
 import requests
 import hashlib
 from typing import List
+from urllib.parse import urlparse as _urlparse
 from .title import Title
 from .tmd import TMD
 from .ticket import Ticket
@@ -13,7 +14,8 @@ from .ticket import Ticket
 _nus_endpoint = ["http://nus.cdn.shop.wii.com/ccs/download/", "http://ccs.cdn.wup.shop.nintendo.net/ccs/download/"]
 
 
-def download_title(title_id: str, title_version: int = None, wiiu_endpoint: bool = False) -> Title:
+def download_title(title_id: str, title_version: int = None, wiiu_endpoint: bool = False,
+                   endpoint_override: str = False) -> Title:
     """
     Download an entire title and all of its contents, then load the downloaded components into a Title object for
     further use. This method is NOT recommended for general use, as it has absolutely no verbosity. It is instead
@@ -23,10 +25,13 @@ def download_title(title_id: str, title_version: int = None, wiiu_endpoint: bool
     ----------
     title_id : str
         The Title ID of the title to download.
-    title_version : int, option
+    title_version : int, optional
         The version of the title to download. Defaults to latest if not set.
-    wiiu_endpoint : bool, option
+    wiiu_endpoint : bool, optional
         Whether the Wii U endpoint for the NUS should be used or not. This increases download speeds. Defaults to False.
+    endpoint_override: str, optional
+        A custom endpoint URL to use instead of the standard Wii or Wii U endpoints. Defaults to no override, and if
+        set entirely overrides the "wiiu_endpoint" parameter.
 
     Returns
     -------
@@ -36,17 +41,18 @@ def download_title(title_id: str, title_version: int = None, wiiu_endpoint: bool
     # First, create the new title.
     title = Title()
     # Download and load the TMD, Ticket, and certs.
-    title.load_tmd(download_tmd(title_id, title_version, wiiu_endpoint))
+    title.load_tmd(download_tmd(title_id, title_version, wiiu_endpoint, endpoint_override))
-    title.load_ticket(download_ticket(title_id, wiiu_endpoint))
+    title.load_ticket(download_ticket(title_id, wiiu_endpoint, endpoint_override))
-    title.wad.set_cert_data(download_cert(wiiu_endpoint))
+    title.wad.set_cert_data(download_cert(wiiu_endpoint, endpoint_override))
     # Download all contents
     title.load_content_records()
-    title.content.content_list = download_contents(title_id, title.tmd, wiiu_endpoint)
+    title.content.content_list = download_contents(title_id, title.tmd, wiiu_endpoint, endpoint_override)
     # Return the completed title.
     return title
 
 
-def download_tmd(title_id: str, title_version: int = None, wiiu_endpoint: bool = False) -> bytes:
+def download_tmd(title_id: str, title_version: int = None, wiiu_endpoint: bool = False,
+                 endpoint_override: str = None) -> bytes:
     """
     Downloads the TMD of the Title specified in the object. Will download the latest version by default, or another
     version if it was manually specified in the object.
@@ -59,6 +65,9 @@ def download_tmd(title_id: str, title_version: int = None, wiiu_endpoint: bool =
         The version of the TMD to download. Defaults to latest if not set.
     wiiu_endpoint : bool, option
         Whether the Wii U endpoint for the NUS should be used or not. This increases download speeds. Defaults to False.
+    endpoint_override: str, optional
+        A custom endpoint URL to use instead of the standard Wii or Wii U endpoints. Defaults to no override, and if
+        set entirely overrides the "wiiu_endpoint" parameter.
 
     Returns
     -------
@@ -67,10 +76,14 @@ def download_tmd(title_id: str, title_version: int = None, wiiu_endpoint: bool =
     """
     # Build the download URL. The structure is download/<TID>/tmd for latest and download/<TID>/tmd.<version> for
     # when a specific version is requested.
-    if wiiu_endpoint is False:
+    if endpoint_override is not None:
-        tmd_url = _nus_endpoint[0] + title_id + "/tmd"
+        endpoint_url = _validate_endpoint(endpoint_override)
     else:
-        tmd_url = _nus_endpoint[1] + title_id + "/tmd"
+        if wiiu_endpoint:
+            endpoint_url = _nus_endpoint[1]
+        else:
+            endpoint_url = _nus_endpoint[0]
+    tmd_url = endpoint_url + title_id + "/tmd"
     # Add the version to the URL if one was specified.
     if title_version is not None:
         tmd_url += "." + str(title_version)
@@ -89,7 +102,7 @@ def download_tmd(title_id: str, title_version: int = None, wiiu_endpoint: bool =
     return tmd
 
 
-def download_ticket(title_id: str, wiiu_endpoint: bool = False) -> bytes:
+def download_ticket(title_id: str, wiiu_endpoint: bool = False, endpoint_override: str = None) -> bytes:
     """
     Downloads the Ticket of the Title specified in the object. This will only work if the Title ID specified is for
     a free title.
@@ -100,6 +113,9 @@ def download_ticket(title_id: str, wiiu_endpoint: bool = False) -> bytes:
         The Title ID of the title to download the Ticket for.
     wiiu_endpoint : bool, option
         Whether the Wii U endpoint for the NUS should be used or not. This increases download speeds. Defaults to False.
+    endpoint_override: str, optional
+        A custom endpoint URL to use instead of the standard Wii or Wii U endpoints. Defaults to no override, and if
+        set entirely overrides the "wiiu_endpoint" parameter.
 
     Returns
     -------
@@ -108,10 +124,14 @@ def download_ticket(title_id: str, wiiu_endpoint: bool = False) -> bytes:
     """
     # Build the download URL. The structure is download/<TID>/cetk, and cetk will only exist if this is a free
     # title.
-    if wiiu_endpoint is False:
+    if endpoint_override is not None:
-        ticket_url = _nus_endpoint[0] + title_id + "/cetk"
+        endpoint_url = _validate_endpoint(endpoint_override)
     else:
-        ticket_url = _nus_endpoint[1] + title_id + "/cetk"
+        if wiiu_endpoint:
+            endpoint_url = _nus_endpoint[1]
+        else:
+            endpoint_url = _nus_endpoint[0]
+    ticket_url = endpoint_url + title_id + "/cetk"
     # Make the request.
     ticket_request = requests.get(url=ticket_url, headers={'User-Agent': 'wii libnup/1.0'}, stream=True)
     if ticket_request.status_code != 200:
@@ -126,7 +146,7 @@ def download_ticket(title_id: str, wiiu_endpoint: bool = False) -> bytes:
     return ticket
 
 
-def download_cert(wiiu_endpoint: bool = False) -> bytes:
+def download_cert(wiiu_endpoint: bool = False, endpoint_override: str = None) -> bytes:
     """
     Downloads the signing certificate used by all WADs. This uses System Menu 4.3U as the source.
 
@@ -134,6 +154,9 @@ def download_cert(wiiu_endpoint: bool = False) -> bytes:
     ----------
     wiiu_endpoint : bool, option
         Whether the Wii U endpoint for the NUS should be used or not. This increases download speeds. Defaults to False.
+    endpoint_override: str, optional
+        A custom endpoint URL to use instead of the standard Wii or Wii U endpoints. Defaults to no override, and if
+        set entirely overrides the "wiiu_endpoint" parameter.
 
     Returns
     -------
@@ -141,12 +164,15 @@ def download_cert(wiiu_endpoint: bool = False) -> bytes:
         The cert file.
     """
     # Download the TMD and cetk for the System Menu 4.3U.
-    if wiiu_endpoint is False:
+    if endpoint_override is not None:
-        tmd_url = _nus_endpoint[0] + "0000000100000002/tmd.513"
+        endpoint_url = _validate_endpoint(endpoint_override)
-        cetk_url = _nus_endpoint[0] + "0000000100000002/cetk"
     else:
-        tmd_url = _nus_endpoint[1] + "0000000100000002/tmd.513"
+        if wiiu_endpoint:
-        cetk_url = _nus_endpoint[1] + "0000000100000002/cetk"
+            endpoint_url = _nus_endpoint[1]
+        else:
+            endpoint_url = _nus_endpoint[0]
+    tmd_url = endpoint_url + "0000000100000002/tmd.513"
+    cetk_url = endpoint_url + "0000000100000002/cetk"
     tmd = requests.get(url=tmd_url, headers={'User-Agent': 'wii libnup/1.0'}, stream=True).content
     cetk = requests.get(url=cetk_url, headers={'User-Agent': 'wii libnup/1.0'}, stream=True).content
     # Assemble the certificate.
@@ -163,7 +189,8 @@ def download_cert(wiiu_endpoint: bool = False) -> bytes:
     return cert
 
 
-def download_content(title_id: str, content_id: int, wiiu_endpoint: bool = False) -> bytes:
+def download_content(title_id: str, content_id: int, wiiu_endpoint: bool = False,
+                     endpoint_override: str = None) -> bytes:
     """
     Downloads a specified content for the title specified in the object.
 
@@ -175,6 +202,9 @@ def download_content(title_id: str, content_id: int, wiiu_endpoint: bool = False
         The Content ID of the content you wish to download.
     wiiu_endpoint : bool, option
         Whether the Wii U endpoint for the NUS should be used or not. This increases download speeds. Defaults to False.
+    endpoint_override: str, optional
+        A custom endpoint URL to use instead of the standard Wii or Wii U endpoints. Defaults to no override, and if
+        set entirely overrides the "wiiu_endpoint" parameter.
 
     Returns
     -------
@@ -185,10 +215,14 @@ def download_content(title_id: str, content_id: int, wiiu_endpoint: bool = False
     content_id_hex = hex(content_id)[2:]
     if len(content_id_hex) < 2:
         content_id_hex = "0" + content_id_hex
-    if wiiu_endpoint is False:
+    if endpoint_override is not None:
-        content_url = _nus_endpoint[0] + title_id + "/000000" + content_id_hex
+        endpoint_url = _validate_endpoint(endpoint_override)
     else:
-        content_url = _nus_endpoint[1] + title_id + "/000000" + content_id_hex
+        if wiiu_endpoint:
+            endpoint_url = _nus_endpoint[1]
+        else:
+            endpoint_url = _nus_endpoint[0]
+    content_url = endpoint_url + title_id + "/000000" + content_id_hex
     # Make the request.
     content_request = requests.get(url=content_url, headers={'User-Agent': 'wii libnup/1.0'}, stream=True)
     if content_request.status_code != 200:
@@ -199,7 +233,8 @@ def download_content(title_id: str, content_id: int, wiiu_endpoint: bool = False
     return content_data
 
 
-def download_contents(title_id: str, tmd: TMD, wiiu_endpoint: bool = False) -> List[bytes]:
+def download_contents(title_id: str, tmd: TMD, wiiu_endpoint: bool = False,
+                      endpoint_override: str = None) -> List[bytes]:
     """
     Downloads all the contents for the title specified in the object. This requires a TMD to already be available
     so that the content records can be accessed.
@@ -212,6 +247,9 @@ def download_contents(title_id: str, tmd: TMD, wiiu_endpoint: bool = False) -> L
         The TMD that matches the title that the contents being downloaded are from.
     wiiu_endpoint : bool, option
         Whether the Wii U endpoint for the NUS should be used or not. This increases download speeds. Defaults to False.
+    endpoint_override: str, optional
+        A custom endpoint URL to use instead of the standard Wii or Wii U endpoints. Defaults to no override, and if
+        set entirely overrides the "wiiu_endpoint" parameter.
 
     Returns
     -------
@@ -228,6 +266,29 @@ def download_contents(title_id: str, tmd: TMD, wiiu_endpoint: bool = False) -> L
     content_list = []
     for content_id in content_ids:
         # Call self.download_content() for each Content ID.
-        content = download_content(title_id, content_id, wiiu_endpoint)
+        content = download_content(title_id, content_id, wiiu_endpoint, endpoint_override)
         content_list.append(content)
     return content_list
+
+
+def _validate_endpoint(endpoint: str) -> str:
+    """
+    Validate the provided NUS endpoint URL and append the required path if necessary.
+
+    Parameters
+    ----------
+    endpoint: str
+        The NUS endpoint URL to validate.
+
+    Returns
+    -------
+    str
+        The validated NUS endpoint with the proper path.
+    """
+    # Find the root of the URL and then assemble the correct URL based on that.
+    new_url = _urlparse(endpoint)
+    if new_url.netloc == "":
+        endpoint_url = "http://" + new_url.path + "/ccs/download/"
+    else:
+        endpoint_url = "http://" + new_url.netloc + "/ccs/download/"
+    return endpoint_url