diff --git a/src/libWiiPy/title/nus.py b/src/libWiiPy/title/nus.py index cbf0fa6..ddfca37 100644 --- a/src/libWiiPy/title/nus.py +++ b/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_ticket(download_ticket(title_id, wiiu_endpoint)) - title.wad.set_cert_data(download_cert(wiiu_endpoint)) + title.load_tmd(download_tmd(title_id, title_version, wiiu_endpoint, endpoint_override)) + title.load_ticket(download_ticket(title_id, wiiu_endpoint, endpoint_override)) + 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//tmd for latest and download//tmd. for # when a specific version is requested. - if wiiu_endpoint is False: - tmd_url = _nus_endpoint[0] + title_id + "/tmd" + if endpoint_override is not None: + 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//cetk, and cetk will only exist if this is a free # title. - if wiiu_endpoint is False: - ticket_url = _nus_endpoint[0] + title_id + "/cetk" + if endpoint_override is not None: + 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: - tmd_url = _nus_endpoint[0] + "0000000100000002/tmd.513" - cetk_url = _nus_endpoint[0] + "0000000100000002/cetk" + if endpoint_override is not None: + endpoint_url = _validate_endpoint(endpoint_override) else: - tmd_url = _nus_endpoint[1] + "0000000100000002/tmd.513" - cetk_url = _nus_endpoint[1] + "0000000100000002/cetk" + if wiiu_endpoint: + 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: - content_url = _nus_endpoint[0] + title_id + "/000000" + content_id_hex + if endpoint_override is not None: + 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