Skip to content

Commit

Permalink
Add support for DMS video delivery
Browse files Browse the repository at this point in the history 
Continues work on #139
  • Loading branch information
@AlexAplin
AlexAplin committed Mar 3, 2024
1 parent 1f6de8a commit 67a2d54
Showing 1 changed file with 114 additions and 41 deletions.
155 changes: 114 additions & 41 deletions nndownload/nndownload.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,12 @@
import sys
import threading
import time
import traceback
import xml.dom.minidom
from typing import AnyStr, List, Match

import aiohttp
import requests
import ffmpeg
from aiohttp_socks import ProxyConnector
from bs4 import BeautifulSoup
from mutagen.mp4 import MP4, MP4StreamInfoError
Expand Down Expand Up @@ -64,6 +64,8 @@
r"(?(6)/((?:[a-z]{2})?\d+))?(?:\?(?:user_id=(.*)|.*)?)?$")
M3U8_STREAM_RE = re.compile(r"(?:(?:#EXT-X-STREAM-INF)|#EXT-X-I-FRAME-STREAM-INF):.*(?:BANDWIDTH=(\d+)).*\n(.*)")
M3U8_MEDIA_RE = re.compile(r"(?:#EXT-X-MEDIA:TYPE=)(?:(\w+))(?:.*),URI=\"(.*)\"")
M3U8_KEY_RE = re.compile(r"((?:#EXT-X-KEY)(?:.*),?URI=\")(.*)\"(.*)")
M3U8_MAP_RE = re.compile(r"((?:#EXT-X-MAP)(?:.*),?URI=\")(.*)\"(.*)")
SEIGA_DRM_KEY_RE = re.compile(r"/image/([a-z0-9]+)")
SEIGA_USER_ID_RE = re.compile(r"user_id=(\d+)")
SEIGA_MANGA_ID_RE = re.compile(r"/comic/(\d+)")
Expand Down Expand Up @@ -408,6 +410,18 @@ def find_extension(mimetype: AnyStr) -> AnyStr:
return MIMETYPES.get(mimetype) or mimetypes.guess_extension(mimetype, strict=True)


def generic_dl_request(session: requests.Session, uri: AnyStr, filename: AnyStr, binary: bool=False):
"""Generic request to download and write to file."""

request = session.get(uri)
request.raise_for_status()
request_body = request.content if binary else request.text
mode = "wb" if binary else "w"
with open(filename, mode) as file:
file.write(request_body)
return request_body


## Nama methods

def generate_stream(session: requests.Session, master_url: AnyStr) -> AnyStr:
Expand Down Expand Up @@ -464,6 +478,7 @@ async def open_nama_websocket(
is_timeshift: bool = False
):
"""Open a WebSocket connection to receive and generate the stream playlist URL."""

proxy = session.proxies.get("http://") # Same mount as https://
connector = ProxyConnector.from_url(proxy) if proxy else None
async with aiohttp.ClientSession(connector=connector) as websocket_session:
Expand Down Expand Up @@ -1026,7 +1041,7 @@ def request_video(session: requests.Session, video_id: AnyStr):
if not _cmdl_opts.skip_media:
continue_code = download_video_media(session, filename, template_params)
if _cmdl_opts.break_on_existing and not continue_code:
raise ExistingDownloadEncountered("Exiting as an existing video was encountered")
raise ExistingDownloadEncounteredQuit("Exiting as an existing video was encountered")
if _cmdl_opts.add_metadata:
add_metadata_to_video(filename, template_params)
if _cmdl_opts.dump_metadata:
Expand Down Expand Up @@ -1203,18 +1218,61 @@ def download_video_part(session: requests.Session, start, end, filename: AnyStr,
update_multithread_progress(len(block))


def perform_ffmpeg_dl(filename: AnyStr, streams: List):
"""Send video and/or audio stream to ffmpeg for download."""

# TODO: Better progress reporting
# TODO: Overwrite detection
# TODO: Catch failures
inputs = []
for stream in streams:
input = ffmpeg.input(stream, protocol_whitelist="https,http,tls,tcp,file,crypto", allowed_extensions="ALL")
inputs.append(input)
output = ffmpeg.output(*inputs, filename, vcodec="copy", acodec="copy")
output.run()


def rewrite_file(filename: AnyStr, old_str: AnyStr, new_str: AnyStr):
"""Replace a string in a text file."""

with open(filename, "r+") as file:
raw = file.read()
new = raw.replace(old_str, new_str)
file.seek(0)
file.write(new)
file.truncate()

def download_video_media(session: requests.Session, filename: AnyStr, template_params: dict):
"""Download video from response URL and display progress."""

output("Downloading {0} to \"{1}\"...\n".format(template_params["id"], filename), logging.INFO)

# Dwango Media Service (DMS)
if template_params.get("video_uri") or template_params.get("audio_uri"):
output("Downloading videos delivered with Dwango Media Service (DMS) is not currently supported.\n", logging.WARNING)
return
cwd = os.path.dirname(os.path.realpath(filename))
m3u8_streams = []
# TODO: Fix clumsy parameters check
for stream_type in ["video_uri", "audio_uri"]:
if template_params.get(stream_type):
# TODO: Consider other extensions
# TODO: Perform cleanup of temp files
m3u8_path = os.path.join(cwd, f"{template_params['id']}_{stream_type}.m3u8.temp")
m3u8 = generic_dl_request(session, template_params[stream_type], m3u8_path)

# It's minimally viable to only rewrite the key file locally for now
# Might be wise to eventually do this with the map and all individual segments
key_match = M3U8_KEY_RE.search(m3u8)
if not key_match:
raise FormatNotAvailableException("Could not retrieve key file from manifest")
key_url = key_match[2]
key_path = os.path.join(cwd, f"{template_params['id']}_{stream_type}.key.temp")
generic_dl_request(session, key_url, key_path, binary=True)
rewrite_file(m3u8_path, key_url, key_path)
m3u8_streams.append(m3u8_path)
perform_ffmpeg_dl(filename, m3u8_streams)
return True

# Dwango Media Cluster (DMC)

dl_stream = session.head(template_params["url"])
dl_stream.raise_for_status()
video_len = int(dl_stream.headers["content-length"])
Expand Down Expand Up @@ -1375,11 +1433,11 @@ def list_qualities(sources_type: str, sources: list, is_dms: bool):
output("{:<24} | {:<10} | {:<46}\n".format(source_id, str(is_available), quality_aggregate), logging.INFO, force=True)


def select_dmc_quality(template_params: dict, template_key: AnyStr, sources: list, quality="") -> List[AnyStr]:
"""Select the specified quality from a sources list on DMC videos."""
def select_quality(template_params: dict, template_key: AnyStr, sources: list, quality="") -> List[AnyStr]:
"""Select the specified quality from a sources list on DMC and DMS videos."""

if quality and _cmdl_opts.force_high_quality:
output("-f/--force-high-quality active. Ignoring quality...\n", logging.WARNING)
output("-f/--force-high-quality was set. Ignoring specified quality...\n", logging.WARNING)

# Assumes qualities are in descending order
highest_quality = sources[0]
Expand Down Expand Up @@ -1440,37 +1498,52 @@ def perform_api_request(session: requests.Session, document: BeautifulSoup) -> d
# Perform request to Dwango Media Service (DMS)
# Began rollout starting 2023-11-01 for select videos and users (https://blog.nicovideo.jp/niconews/205042.html)
# Videos longer than 30 minutes in HD (>720p) quality appear to be served this way exclusively
# elif params["media"]["domand"]:
# if _cmdl_opts.list_qualities:
# list_qualities("video", params["media"]["domand"]["videos"], True)
# list_qualities("audio", params["media"]["domand"]["audios"], True)
# raise ListQualitiesQuit("Exiting after listing available qualities")

# video_id = params["video"]["id"]
# access_right_key = params["media"]["domand"]["accessRightKey"]
# watch_track_id = params["client"]["watchTrackId"]

# # Limited to one video and audio source
# payload = json.dumps({"outputs":[["video-h264-720p","audio-aac-128kbps",]]})

# output("Retrieving video manifest ...\n", logging.INFO)
# headers = {
# "X-Access-Right-Key": access_right_key,
# "X-Request-With": "https://www.nicovideo.jp", # Only provided on this endpoint
# }
# session.options(VIDEO_DMS_WATCH_API.format(video_id, watch_track_id)) # OPTIONS
# get_manifest_request = session.post(VIDEO_DMS_WATCH_API.format(video_id, watch_track_id), headers={**API_HEADERS, **headers}, data=payload)
# get_manifest_request.raise_for_status()
# manifest_url = get_manifest_request.json()["data"]["contentUrl"]
# manifest_request = session.get(manifest_url)
# manifest_request.raise_for_status()
# manifest_text = manifest_request.text
# output("Retrieved video manifest.\n", logging.INFO)

# output("Collecting video media URIs...\n")
# template_params["video_uri"] = get_stream_from_manifest(manifest_text)
# template_params["audio_uri"] = audio_playlist = get_media_from_manifest(manifest_text, "audio")
# output("Collected video media URIs.\n", logging.INFO)
elif params["media"]["domand"]:
if _cmdl_opts.list_qualities:
list_qualities("video", params["media"]["domand"]["videos"], True)
list_qualities("audio", params["media"]["domand"]["audios"], True)
raise ListQualitiesQuit("Exiting after listing available qualities")

video_id = params["video"]["id"]
access_right_key = params["media"]["domand"]["accessRightKey"]
watch_track_id = params["client"]["watchTrackId"]

video_sources = select_quality(
template_params,
"video_quality",
params["media"]["domand"]["videos"],
_cmdl_opts.video_quality or "highest"
)
audio_sources = select_quality(
template_params,
"audio_quality",
params["media"]["domand"]["audios"],
_cmdl_opts.audio_quality or "highest"
)

# Limited to one video and audio source
video_source = video_sources[0]
audio_source = audio_sources[0]
payload = json.dumps({"outputs":[[video_source,audio_source,]]})

output("Retrieving video manifest ...\n", logging.INFO)
headers = {
"X-Access-Right-Key": access_right_key,
"X-Request-With": "https://www.nicovideo.jp", # Only provided on this endpoint
}
session.options(VIDEO_DMS_WATCH_API.format(video_id, watch_track_id)) # OPTIONS
get_manifest_request = session.post(VIDEO_DMS_WATCH_API.format(video_id, watch_track_id), headers={**API_HEADERS, **headers}, data=payload)
get_manifest_request.raise_for_status()
manifest_url = get_manifest_request.json()["data"]["contentUrl"]
manifest_request = session.get(manifest_url)
manifest_request.raise_for_status()
manifest_text = manifest_request.text
output("Retrieved video manifest.\n", logging.INFO)

output("Collecting video media URIs...\n")
template_params["video_uri"] = get_stream_from_manifest(manifest_text)
template_params["audio_uri"] = get_media_from_manifest(manifest_text, "audio")
output("Collected video media URIs.\n", logging.INFO)

# Perform request to Dwango Media Cluster (DMC)
elif params["media"]["delivery"]:
Expand All @@ -1489,13 +1562,13 @@ def perform_api_request(session: requests.Session, document: BeautifulSoup) -> d
file_extension = template_params["ext"]
priority = params["media"]["delivery"]["movie"]["session"]["priority"]

video_sources = select_dmc_quality(
video_sources = select_quality(
template_params,
"video_quality",
params["media"]["delivery"]["movie"]["videos"],
_cmdl_opts.video_quality
)
audio_sources = select_dmc_quality(
audio_sources = select_quality(
template_params,
"audio_quality",
params["media"]["delivery"]["movie"]["audios"],
Expand Down

0 comments on commit 67a2d54

Please sign in to comment.