From c1fbdfadea3495d9bd386d3d755078681fe232b8 Mon Sep 17 00:00:00 2001
From: Nick Dowsett <nickdowsett42@gmail.com>
Date: Mon, 7 Oct 2024 23:05:21 +0800
Subject: [PATCH] Add ability to supply own PO tokens

---
 src/constants.rs |  2 --
 src/info.rs      | 19 ++++++++++++++++---
 src/structs.rs   |  3 +++
 src/utils.rs     | 25 +------------------------
 4 files changed, 20 insertions(+), 29 deletions(-)

diff --git a/src/constants.rs b/src/constants.rs
index 6a0d0e4..9a5e551 100644
--- a/src/constants.rs
+++ b/src/constants.rs
@@ -52,8 +52,6 @@ pub(crate) const DEFAULT_DL_CHUNK_SIZE: u64 = 10485760;
 /// Default max number of retries for a web reqwest.
 pub(crate) const DEFAULT_MAX_RETRIES: u32 = 3;
 
-pub(crate) const POTOKEN_EXPERIMENTS: &[&str] = &["51217476", "51217102"];
-
 pub static INNERTUBE_CLIENT: Lazy<HashMap<&str, (&str, &str, &str)>> =
     // (clientVersion, clientName, json value)
     Lazy::new(|| {
diff --git a/src/info.rs b/src/info.rs
index 9e46521..15a7e4e 100644
--- a/src/info.rs
+++ b/src/info.rs
@@ -6,6 +6,7 @@ use reqwest::{
 use reqwest_middleware::{ClientBuilder, ClientWithMiddleware};
 use reqwest_retry::{policies::ExponentialBackoff, RetryTransientMiddleware};
 use scraper::{Html, Selector};
+use serde_json::json;
 use std::{borrow::{Borrow, Cow}, path::Path, time::Duration};
 use url::Url;
 
@@ -22,7 +23,7 @@ use crate::{
         CustomRetryableStrategy, PlayerResponse, VideoError, VideoInfo, VideoOptions, YTConfig,
     },
     utils::{
-        between, check_experiments, choose_format, clean_video_details, get_functions, get_html,
+        between, choose_format, clean_video_details, get_functions, get_html,
         get_html5player, get_random_v6_ip, get_video_id, get_ytconfig, is_age_restricted_from_html,
         is_live, is_not_yet_broadcasted, is_play_error, is_player_response_error, is_private_video,
         is_rental, parse_live_video_formats, parse_video_formats, sort_formats,
@@ -188,11 +189,12 @@ impl<'opts> Video<'opts> {
         }
 
         // POToken experiment detected fallback to ios client (Webpage contains broken formats)
-        if check_experiments(&response) && !is_live(&player_response) {
+        if !is_live(&player_response) {
             let ios_ytconfig = self
                 .get_player_ytconfig(
                     &response,
                     INNERTUBE_CLIENT.get("ios").cloned().unwrap_or_default(),
+                    self.options.request_options.po_token.as_ref()
                 )
                 .await?;
 
@@ -210,6 +212,7 @@ impl<'opts> Video<'opts> {
                         .get("tv_embedded")
                         .cloned()
                         .unwrap_or_default(),
+                    self.options.request_options.po_token.as_ref()
                 )
                 .await?;
 
@@ -521,6 +524,7 @@ impl<'opts> Video<'opts> {
         &self,
         html: &str,
         configs: (&str, &str, &str),
+        po_token: Option<&String>,
     ) -> Result<String, VideoError> {
         use std::str::FromStr;
 
@@ -530,7 +534,7 @@ impl<'opts> Video<'opts> {
         let sts = ytcfg.sts.unwrap_or(0);
         let video_id = self.get_video_id();
 
-        let query = serde_json::from_str::<serde_json::Value>(&format!(
+        let mut query = serde_json::from_str::<serde_json::Value>(&format!(
             r#"{{
             {client}
             "playbackContext": {{
@@ -543,6 +547,15 @@ impl<'opts> Video<'opts> {
         }}"#
         ))
         .unwrap_or_default();
+        if let Some(po_token) = po_token {
+            query
+                .as_object_mut()
+                .expect("Declared as object above")
+                .insert(
+                    "serviceIntegrityDimensions".to_string(),
+                    json!({"poToken": po_token})
+                );
+        }
 
         static CONFIGS: Lazy<(HeaderMap, &str)> = Lazy::new(|| {
             (HeaderMap::from_iter([
diff --git a/src/structs.rs b/src/structs.rs
index 5a9340e..b84bf94 100644
--- a/src/structs.rs
+++ b/src/structs.rs
@@ -248,6 +248,9 @@ pub struct RequestOptions {
     ///     };
     /// ```
     pub max_retries: Option<u32>,
+    /// Supply a YouTube Proof of Origin token. Use at your own risk.
+    /// See https://github.com/yt-dlp/yt-dlp/wiki/Extractors#po-token-guide for more information.
+    pub po_token: Option<String>,
 }
 
 #[derive(thiserror::Error, Debug)]
diff --git a/src/utils.rs b/src/utils.rs
index dcbb062..43b9ec1 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -16,7 +16,7 @@ use urlencoding::decode;
 use crate::{
     constants::{
         AGE_RESTRICTED_URLS, AUDIO_ENCODING_RANKS, BASE_URL, FORMATS, IPV6_REGEX, PARSE_INT_REGEX,
-        POTOKEN_EXPERIMENTS, VALID_QUERY_DOMAINS, VIDEO_ENCODING_RANKS,
+        VALID_QUERY_DOMAINS, VIDEO_ENCODING_RANKS,
     },
     info_extras::{get_author, get_chapters, get_dislikes, get_likes, get_storyboards},
     structs::{
@@ -996,29 +996,6 @@ pub fn get_ytconfig(html: &str) -> Result<YTConfig, VideoError> {
     }
 }
 
-pub fn check_experiments(html: &str) -> bool {
-    if let Some(configs) = get_ytconfig(html)
-        .ok()
-        .and_then(|x| x.web_player_context_configs.clone())
-        .and_then(|v| v.as_object().cloned())
-    {
-        return configs.iter().any(|(_, config)| {
-            config
-                .get("serializedExperimentIds")
-                .and_then(|v| v.as_str())
-                .map(|ids| {
-                    let ids_set = ids.split(',').collect::<Vec<&str>>();
-                    POTOKEN_EXPERIMENTS
-                        .iter()
-                        .any(|token| ids_set.contains(token))
-                })
-                .unwrap_or(false)
-        });
-    }
-
-    false
-}
-
 type CacheFunctions = Lazy<RwLock<Option<(String, Vec<(String, String)>)>>>;
 static FUNCTIONS: CacheFunctions = Lazy::new(|| RwLock::new(None));