path: root/player/lua/ytdl_hook.lua

local utils = require 'mp.utils'
local msg = require 'mp.msg'
local options = require 'mp.options'

local o = {
    exclude = "",
    try_ytdl_first = false,
    use_manifests = false,
    all_formats = false,
    force_all_formats = true,
    thumbnails = "none",
    ytdl_path = "",
}

local ytdl = {
    path = "",
    paths_to_search = {"yt-dlp", "yt-dlp_x86", "youtube-dl"},
    searched = false,
    blacklisted = {}
}

options.read_options(o, nil, function()
    ytdl.blacklisted = {} -- reparse o.exclude next time
    ytdl.searched = false
end)

local chapter_list = {}
local playlist_cookies = {}

function Set (t)
    local set = {}
    for _, v in pairs(t) do set[v] = true end
    return set
end

-- ?: surrogate (keep in mind that there is no lazy evaluation)
function iif(cond, if_true, if_false)
    if cond then
        return if_true
    end
    return if_false
end

-- youtube-dl JSON name to mpv tag name
local tag_list = {
    ["uploader"]        = "uploader",
    ["channel_url"]     = "channel_url",
    -- these titles tend to be a bit too long, so hide them on the terminal
    -- (default --display-tags does not include this name)
    ["description"]     = "ytdl_description",
    -- "title" is handled by force-media-title
    -- tags don't work with all_formats=yes
}

local safe_protos = Set {
    "http", "https", "ftp", "ftps",
    "rtmp", "rtmps", "rtmpe", "rtmpt", "rtmpts", "rtmpte",
    "data"
}

-- For some sites, youtube-dl returns the audio codec (?) only in the "ext" field.
local ext_map = {
    ["mp3"]         = "mp3",
    ["opus"]        = "opus",
}

local codec_map = {
    -- src pattern  = mpv codec
    ["vtt"]         = "webvtt",
    ["opus"]        = "opus",
    ["vp9"]         = "vp9",
    ["avc1%..*"]    = "h264",
    ["av01%..*"]    = "av1",
    ["mp4a%..*"]    = "aac",
}

-- Codec name as reported by youtube-dl mapped to mpv internal codec names.
-- Fun fact: mpv will not really use the codec, but will still try to initialize
-- the codec on track selection (just to scrap it), meaning it's only a hint,
-- but one that may make initialization fail. On the other hand, if the codec
-- is valid but completely different from the actual media, nothing bad happens.
local function map_codec_to_mpv(codec)
    if codec == nil then
        return nil
    end
    for k, v in pairs(codec_map) do
        local s, e = codec:find(k)
        if s == 1 and e == #codec then
            return v
        end
    end
    return nil
end

local function platform_is_windows()
    return mp.get_property_native("platform") == "windows"
end

local function exec(args)
    msg.debug("Running: " .. table.concat(args, " "))

    return mp.command_native({
        name = "subprocess",
        args = args,
        capture_stdout = true,
        capture_stderr = true,
    })
end

-- return true if it was explicitly set on the command line
local function option_was_set(name)
    return mp.get_property_bool("option-info/" ..name.. "/set-from-commandline",
                                false)
end

-- return true if the option was set locally
local function option_was_set_locally(name)
    return mp.get_property_bool("option-info/" ..name.. "/set-locally", false)
end

-- youtube-dl may set special http headers for some sites (user-agent, cookies)
local function set_http_headers(http_headers)
    if not http_headers then
        return
    end
    local headers = {}
    local useragent = http_headers["User-Agent"]
    if useragent and not option_was_set("user-agent") then
        mp.set_property("file-local-options/user-agent", useragent)
    end
    local additional_fields = {"Cookie", "Referer", "X-Forwarded-For"}
    for idx, item in pairs(additional_fields) do
        local field_value = http_headers[item]
        if field_value then
            headers[#headers + 1] = item .. ": " .. field_value
        end
    end
    if #headers > 0 and not option_was_set("http-header-fields") then
        mp.set_property_native("file-local-options/http-header-fields", headers)
    end
end

local special_cookie_field_names = Set {
    "expires", "max-age", "domain", "path"
}

-- parse single-line Set-Cookie syntax
local function parse_cookies(cookies_line)
    if not cookies_line then
        return {}
    end
    local cookies = {}
    local cookie = {}
    for stem in cookies_line:gmatch('[^;]+') do
        stem = stem:gsub("^%s*(.-)%s*$", "%1")
        local name, value = stem:match('^(.-)=(.+)$')
        if name and name ~= "" and value then
            local cmp_name = name:lower()
            if special_cookie_field_names[cmp_name] then
                cookie[cmp_name] = value
            else
                if cookie.name and cookie.value then
                    table.insert(cookies, cookie)
                end
                cookie = {
                    name = name,
                    value = value,
                }
            end
        end
    end
    if cookie.name and cookie.value then
        local cookie_key = cookie.domain .. ":" .. cookie.name
        cookies[cookie_key] = cookie
    end
    return cookies
end

-- serialize cookies for avformat
local function serialize_cookies_for_avformat(cookies)
    local result = ''
    for _, cookie in pairs(cookies) do
        local cookie_str = ('%s=%s; '):format(cookie.name, cookie.value)
        for k, v in pairs(cookie) do
            if k ~= "name" and k ~= "value" then
                cookie_str = cookie_str .. ('%s=%s; '):format(k, v)
            end
        end
        result = result .. cookie_str .. '\r\n'
    end
    return result
end

-- set file-local cookies, preserving existing ones
local function set_cookies(cookies)
    if not cookies or cookies == "" then
        return
    end

    local option_key = "file-local-options/stream-lavf-o"
    local stream_opts = mp.get_property_native(option_key, {})
    local existing_cookies = parse_cookies(stream_opts["cookies"])

    local new_cookies = parse_cookies(cookies)
    for cookie_key, cookie in pairs(new_cookies) do
        if not existing_cookies[cookie_key] then
            existing_cookies[cookie_key] = cookie
        end
    end

    stream_opts["cookies"] = serialize_cookies_for_avformat(existing_cookies)
    mp.set_property_native(option_key, stream_opts)
end

local function append_libav_opt(props, name, value)
    if not props then
        props = {}
    end

    if name and value and not props[name] then
        props[name] = value
    end

    return props
end

local function edl_escape(url)
    return "%" .. string.len(url) .. "%" .. url
end

local function url_is_safe(url)
    local proto = type(url) == "string" and url:match("^(%a[%w+.-]*):") or nil
    local safe = proto and safe_protos[proto]
    if not safe then
        msg.error(("Ignoring potentially unsafe url: '%s'"):format(url))
    end
    return safe
end

local function time_to_secs(time_string)
    local ret

    local a, b, c = time_string:match("(%d+):(%d%d?):(%d%d)")
    if a ~= nil then
        ret = (a*3600 + b*60 + c)
    else
        a, b = time_string:match("(%d%d?):(%d%d)")
        if a ~= nil then
            ret = (a*60 + b)
        end
    end

    return ret
end

local function extract_chapters(data, video_length)
    local ret = {}

    for line in data:gmatch("[^\r\n]+") do
        local time = time_to_secs(line)
        if time and (time < video_length) then
            table.insert(ret, {time = time, title = line})
        end
    end
    table.sort(ret, function(a, b) return a.time < b.time end)
    return ret
end

local function is_blacklisted(url)
    if o.exclude == "" then return false end
    if #ytdl.blacklisted == 0 then
        for match in o.exclude:gmatch('%|?([^|]+)') do
            ytdl.blacklisted[#ytdl.blacklisted + 1] = match
        end
    end
    if #ytdl.blacklisted > 0 then
        url = url:match('https?://(.+)')
        for _, exclude in ipairs(ytdl.blacklisted) do
            if url:match(exclude) then
                msg.verbose('URL matches excluded substring. Skipping.')
                return true
            end
        end
    end
    return false
end

local function parse_yt_playlist(url, json)
    -- return 0-based index to use with --playlist-start

    if not json.extractor or
       (json.extractor ~= "youtube:tab" and
        json.extractor ~= "youtube:playlist") then
        return nil
    end

    local query = url:match("%?.+")
    if not query then return nil end

    local args = {}
    for arg, param in query:gmatch("(%a+)=([^&?]+)") do
        if arg and param then
            args[arg] = param
        end
    end

    local maybe_idx = tonumber(args["index"])

    -- if index matches v param it's probably the requested item
    if maybe_idx and #json.entries >= maybe_idx and
        json.entries[maybe_idx].id == args["v"] then
        msg.debug("index matches requested video")
        return maybe_idx - 1
    end

    -- if there's no index or it doesn't match, look for video
    for i = 1, #json.entries do
        if json.entries[i].id == args["v"] then
            msg.debug("found requested video in index " .. (i - 1))
            return i - 1
        end
    end

    msg.debug("requested video not found in playlist")
    -- if item isn't on the playlist, give up
    return nil
end

local function make_absolute_url(base_url, url)
    if url:find("https?://") == 1 then return url end

    local proto, domain, rest =
        base_url:match("(https?://)([^/]+/)(.*)/?")
    local segs = {}
    rest:gsub("([^/]+)", function(c) table.insert(segs, c) end)
    url:gsub("([^/]+)", function(c) table.insert(segs, c) end)
    local resolved_url = {}
    for i, v in ipairs(segs) do
        if v == ".." then
            table.remove(resolved_url)
        elseif v ~= "." then
            table.insert(resolved_url, v)
        end
    end
    return proto .. domain ..
        table.concat(resolved_url, "/")
end

local function join_url(base_url, fragment)
    local res = ""
    if base_url and fragment.path then
        res = make_absolute_url(base_url, fragment.path)
    elseif fragment.url then
        res = fragment.url
    end
    return res
end

local function edl_track_joined(fragments, protocol, is_live, base)
    if type(fragments) ~= "table" or not fragments[1] then
        msg.debug("No fragments to join into EDL")
        return nil
    end

    local edl = "edl://"
    local offset = 1
    local parts = {}

    if protocol == "http_dash_segments" and not is_live then
        msg.debug("Using dash")
        local args = ""

        -- assume MP4 DASH initialization segment
        if not fragments[1].duration and #fragments > 1 then
            msg.debug("Using init segment")
            args = args .. ",init=" .. edl_escape(join_url(base, fragments[1]))
            offset = 2
        end

        table.insert(parts, "!mp4_dash" .. args)

        -- Check remaining fragments for duration;
        -- if not available in all, give up.
        for i = offset, #fragments do
            if not fragments[i].duration then
                msg.verbose("EDL doesn't support fragments " ..
                         "without duration with MP4 DASH")
                return nil
            end
        end
    end

    for i = offset, #fragments do
        local fragment = fragments[i]
        if not url_is_safe(join_url(base, fragment)) then
            return nil
        end
        table.insert(parts, edl_escape(join_url(base, fragment)))
        if fragment.duration then
            parts[#parts] =
                parts[#parts] .. ",length="..fragment.duration
        end
    end
    return edl .. table.concat(parts, ";") .. ";"
end

local function has_native_dash_demuxer()
    local demuxers = mp.get_property_native("demuxer-lavf-list", {})
    for _, v in ipairs(demuxers) do
        if v == "dash" then
            return true
        end
    end
    return false
end

local function valid_manifest(json)
    local reqfmt = json["requested_formats"] and json["requested_formats"][1] or {}
    if not reqfmt["manifest_url"] and not json["manifest_url"] then
        return false
    end
    local proto = reqfmt["protocol"] or json["protocol"] or ""
    return (proto == "http_dash_segments" and has_native_dash_demuxer()) or
        proto:find("^m3u8")
end

local function as_integer(v, def)
    def = def or 0
    local num = math.floor(tonumber(v) or def)
    if num > -math.huge and num < math.huge then
        return num
    end
    return def
end

local function tags_to_edl(json)
    local tags = {}
    for json_name, mp_name in pairs(tag_list) do
        local v = json[json_name]
        if v then
            tags[#tags + 1] = mp_name .. "=" .. edl_escape(tostring(v))
        end
    end
    if #tags == 0 then
        return nil
    end
    return "!global_tags," .. table.concat(tags, ",")
end

-- Convert a format list from youtube-dl to an EDL URL, or plain URL.
--  json: full json blob by youtube-dl
--  formats: format list by youtube-dl
--  use_all_formats: if=true, then formats is the full format list, and the
--                   function will attempt to return them as delay-loaded tracks
-- See res table initialization in the function for result type.
local function formats_to_edl(json, formats, use_all_formats)
    local res = {
        -- the media URL, which may be EDL
        url = nil,
        -- for use_all_formats=true: whether any muxed formats are present, and
        -- at the same time the separate EDL parts don't have both audio/video
        muxed_needed = false,
    }

    local default_formats = {}
    local requested_formats = json["requested_formats"] or json["requested_downloads"]
    if use_all_formats and requested_formats then
        for _, track in ipairs(requested_formats) do
            local id = track["format_id"]
            if id then
                default_formats[id] = true
            end
        end
    end

    local duration = as_integer(json["duration"])
    local single_url = nil
    local streams = {}

    local tbr_only = true
    for index, track in ipairs(formats) do
        tbr_only = tbr_only and track["tbr"] and
                   (not track["abr"]) and (not track["vbr"])
    end

    local has_requested_video = false
    local has_requested_audio = false
    -- Web players with quality selection always show the highest quality
    -- option at the top. Since tracks are usually listed with the first
    -- track at the top, that should also be the highest quality track.
    -- yt-dlp/youtube-dl sorts it's formats from worst to best.
    -- Iterate in reverse to get best track first.
    for index = #formats, 1, -1 do
        local track = formats[index]
        local edl_track = nil
        edl_track = edl_track_joined(track.fragments,
            track.protocol, json.is_live,
            track.fragment_base_url)
        if not edl_track and not url_is_safe(track.url) then
            msg.error("No safe URL or supported fragmented stream available")
            return nil
        end

        local is_default = default_formats[track["format_id"]]
        local tracks = {}
        -- "none" means it is not a video
        -- nil means it is unknown
        if (o.force_all_formats or track.vcodec) and track.vcodec ~= "none" then
            tracks[#tracks + 1] = {
                media_type = "video",
                codec = map_codec_to_mpv(track.vcodec),
            }
            if is_default then
                has_requested_video = true
            end
        end
        if (o.force_all_formats or track.acodec) and track.acodec ~= "none" then
            tracks[#tracks + 1] = {
                media_type = "audio",
                codec = map_codec_to_mpv(track.acodec) or
                        ext_map[track.ext],
            }
            if is_default then
                has_requested_audio = true
            end
        end

        local url = edl_track or track.url
        local hdr = {"!new_stream", "!no_clip", "!no_chapters"}
        local skip = #tracks == 0
        local params = ""

        if use_all_formats then
            for _, sub in ipairs(tracks) do
                -- A single track that is either audio or video. Delay load it.
                local props = ""
                if sub.media_type == "video" then
                    props = props .. ",w=" .. as_integer(track.width)
                                  .. ",h=" .. as_integer(track.height)
                                  .. ",fps=" .. as_integer(track.fps)
                elseif sub.media_type == "audio" then
                    props = props .. ",samplerate=" .. as_integer(track.asr)
                end
                hdr[#hdr + 1] = "!delay_open,media_type=" .. sub.media_type ..
                    ",codec=" .. (sub.codec or "null") .. props

                -- Add bitrate information etc. for better user selection.
                local byterate = 0
                local rates = {"tbr", "vbr", "abr"}
                if #tracks > 1 then
                    rates = {({video = "vbr", audio = "abr"})[sub.media_type]}
                end
                if tbr_only then
                    rates = {"tbr"}
                end
                for _, f in ipairs(rates) do
                    local br = as_integer(track[f])
                    if br > 0 then
                        byterate = math.floor(br * 1000 / 8)
                        break
                    end
                end
                local title = track.format or track.format_note or ""
                if #tracks > 1 then
                    if #title > 0 then
                        title = title .. " "
                    end
                    title = title .. "muxed-" .. index
                end
                local flags = {}
                if is_default then
                    flags[#flags + 1] = "default"
                end
                hdr[#hdr + 1] = "!track_meta,title=" ..
                    edl_escape(title) .. ",byterate=" .. byterate ..
                    iif(#flags > 0, ",flags=" .. table.concat(flags, "+"), "")
            end

            if duration > 0 then
                params = params .. ",length=" .. duration
            end
        end

        if not skip then
            hdr[#hdr + 1] = edl_escape(url) .. params

            streams[#streams + 1] = table.concat(hdr, ";")
            -- In case there is only 1 of these streams.
            -- Note: assumes it has no important EDL headers
            single_url = url
        end
    end

    local tags = tags_to_edl(json)

    -- Merge all tracks into a single virtual file, but avoid EDL if it's
    -- only a single track without metadata (i.e. redundant).
    if #streams == 1 and single_url and not tags then
        res.url = single_url
    elseif #streams > 0 then
        if tags then
            -- not a stream; just for the sake of concatenating the EDL string
            streams[#streams + 1] = tags
        end
        res.url = "edl://" .. table.concat(streams, ";")
    else
        return nil
    end

    if has_requested_audio ~= has_requested_video then
        local not_req_prop = has_requested_video and "aid" or "vid"
        if mp.get_property(not_req_prop) == "auto" then
            mp.set_property("file-local-options/" .. not_req_prop, "no")
        end
    end

    return res
end

local function add_single_video(json)
    local streamurl = ""
    local format_info = ""
    local max_bitrate = 0
    local requested_formats = json["requested_formats"] or json["requested_downloads"]
    local all_formats = json["formats"]
    local has_requested_formats = requested_formats and #requested_formats > 0
    local http_headers = has_requested_formats
                         and requested_formats[1].http_headers
                         or json.http_headers
    local cookies = has_requested_formats
                    and requested_formats[1].cookies
                    or json.cookies

    if o.use_manifests and valid_manifest(json) then
        -- prefer manifest_url if present
        format_info = "manifest"

        local mpd_url = requested_formats and