From db6ece6c61f1cd7952ae85fd2f3ee48b2b83339c Mon Sep 17 00:00:00 2001 From: ThetaDev Date: Sun, 31 Jul 2022 20:13:18 +0200 Subject: [PATCH] implemented cache --- .gitignore | 2 + Cargo.toml | 1 + src/cache.rs | 174 ++++++++++++++++++++++++++++++++++++++----- src/client/mod.rs | 137 +++++++++++++++------------------- src/client/player.rs | 10 ++- src/deobfuscate.rs | 95 +++++++++-------------- 6 files changed, 261 insertions(+), 158 deletions(-) diff --git a/.gitignore b/.gitignore index 4fffb2f..23e8123 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,4 @@ /target /Cargo.lock + +rusty-tube.json diff --git a/Cargo.toml b/Cargo.toml index 1e3d1fe..a3ed500 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,3 +26,4 @@ chrono = {version = "0.4.19", features = ["serde"]} env_logger = "0.9.0" test-log = "0.2.11" rstest = "0.15.0" +temp_testdir = "0.2.3" diff --git a/src/cache.rs b/src/cache.rs index f38deec..fb6b02d 100644 --- a/src/cache.rs +++ b/src/cache.rs @@ -1,17 +1,25 @@ -use std::{future::Future, sync::Arc}; +use std::{ + fs::File, + future::Future, + io::BufReader, + path::{Path, PathBuf}, + sync::Arc, +}; use anyhow::Result; use chrono::{DateTime, Duration, Utc}; +use log::{error, info}; use serde::{Deserialize, Serialize}; use tokio::sync::Mutex; -#[derive(Default, Debug)] +#[derive(Default, Debug, Clone)] pub struct Cache { + file: Option, data: Arc>, } #[derive(Default, Debug, Clone, Serialize, Deserialize)] -pub struct CacheData { +struct CacheData { desktop_client: Option>, deobf: Option>, } @@ -38,10 +46,10 @@ pub struct DesktopClientData { #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct DeobfData { - js_url: String, - sig_fn: String, - nsig_fn: String, - sts: String, + pub js_url: String, + pub sig_fn: String, + pub nsig_fn: String, + pub sts: String, } impl Cache { @@ -52,10 +60,11 @@ impl Cache { let mut cache = self.data.lock().await; if cache.desktop_client.is_none() - || cache.desktop_client.as_ref().unwrap().last_update < Utc::now() - Duration::days(1) + || cache.desktop_client.as_ref().unwrap().last_update < Utc::now() - Duration::hours(24) { let cdata = updater.await?; cache.desktop_client = Some(CacheEntry::from(cdata.clone())); + self.save(&cache); Ok(cdata) } else { Ok(cache.desktop_client.as_ref().unwrap().data.clone()) @@ -68,33 +77,102 @@ impl Cache { { let mut cache = self.data.lock().await; if cache.deobf.is_none() - || cache.deobf.as_ref().unwrap().last_update < Utc::now() - Duration::days(1) + || cache.deobf.as_ref().unwrap().last_update < Utc::now() - Duration::hours(1) { let deobf_data = updater.await?; cache.deobf = Some(CacheEntry::from(deobf_data.clone())); + self.save(&cache); Ok(deobf_data) } else { Ok(cache.deobf.as_ref().unwrap().data.clone()) } } - pub async fn to_json(&self) -> Result { + pub async fn to_json(&self) -> Result { let cache = self.data.lock().await; - serde_json::to_string(&cache.clone()) + Ok(serde_json::to_string(&cache.clone())?) } - pub async fn from_json(&self, json: &str) -> Result<(), serde_json::Error> { - let cd = serde_json::from_str::(json)?; + pub async fn to_json_file>(&self, path: P) -> Result<()> { + let cache = self.data.lock().await; + Ok(serde_json::to_writer(&File::create(path)?, &cache.clone())?) + } - let mut cache = self.data.lock().await; - *cache = cd; + pub fn from_json(json: &str) -> Self { + let data: CacheData = match serde_json::from_str(json) { + Ok(cd) => cd, + Err(e) => { + error!( + "Could not load cache from json, falling back to default. Error: {}", + e + ); + CacheData::default() + } + }; + Cache { + data: Arc::new(Mutex::new(data)), + file: None, + } + } - Ok(()) + pub fn from_json_file>(path: P) -> Self { + let file = match File::open(path.as_ref()) { + Ok(file) => file, + Err(e) => { + if e.kind() == std::io::ErrorKind::NotFound { + info!( + "Cache json file at {} not found, will be created", + path.as_ref().to_string_lossy() + ) + } else { + error!( + "Could not open cache json file, falling back to default. Error: {}", + e + ); + } + return Cache { + file: Some(path.as_ref().to_path_buf()), + ..Default::default() + }; + } + }; + let data: CacheData = match serde_json::from_reader(BufReader::new(file)) { + Ok(data) => data, + Err(e) => { + error!( + "Could not load cache from json, falling back to default. Error: {}", + e + ); + return Cache { + file: Some(path.as_ref().to_path_buf()), + ..Default::default() + }; + } + }; + Cache { + data: Arc::new(Mutex::new(data)), + file: Some(path.as_ref().to_path_buf()), + } + } + + fn save(&self, cache: &CacheData) { + match self.file.as_ref() { + Some(file) => match File::create(file) { + Ok(file) => match serde_json::to_writer(file, cache) { + Ok(_) => {} + Err(e) => error!("Could not write cache to json. Error: {}", e), + }, + Err(e) => error!("Could not open cache json file. Error: {}", e), + }, + None => {} + } } } #[cfg(test)] mod tests { + use temp_testdir::TempDir; + use super::*; #[tokio::test] @@ -143,9 +221,69 @@ mod tests { ); let json = cache.to_json().await.unwrap(); + let new_cache = Cache::from_json(&json); - let new_cache = Cache::default(); - new_cache.from_json(&json).await.unwrap(); + assert_eq!( + new_cache + .get_desktop_client_data(async { + Ok(DesktopClientData { + client_version: "".to_owned(), + }) + }) + .await + .unwrap(), + cdata + ); + + assert_eq!( + new_cache + .get_deobf_data(async { + Ok(DeobfData { + js_url: "".to_owned(), + nsig_fn: "".to_owned(), + sig_fn: "".to_owned(), + sts: "".to_owned(), + }) + }) + .await + .unwrap(), + deobf_data + ); + } + + #[tokio::test] + async fn test_file() { + let temp = TempDir::default(); + let mut file_path = PathBuf::from(temp.as_ref()); + file_path.push("cache.json"); + + let cache = Cache::from_json_file(file_path.clone()); + + let cdata = cache + .get_desktop_client_data(async { + Ok(DesktopClientData { + client_version: "1.2.3".to_owned(), + }) + }) + .await + .unwrap(); + + let deobf_data = cache + .get_deobf_data(async { + Ok(DeobfData { + js_url: + "https://www.youtube.com/s/player/011af516/player_ias.vflset/en_US/base.js" + .to_owned(), + sig_fn: "t_sig_fn".to_owned(), + nsig_fn: "t_nsig_fn".to_owned(), + sts: "t_sts".to_owned(), + }) + }) + .await + .unwrap(); + + assert!(file_path.exists()); + let new_cache = Cache::from_json_file(file_path.clone()); assert_eq!( new_cache diff --git a/src/client/mod.rs b/src/client/mod.rs index 2f6e277..e8cfb26 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -3,7 +3,7 @@ mod response; use std::{sync::Arc, time::Instant}; -use anyhow::{anyhow, bail, Context, Result}; +use anyhow::{anyhow, Context, Result}; use async_trait::async_trait; use fancy_regex::Regex; use log::{debug, warn}; @@ -13,7 +13,11 @@ use reqwest::{header, Client, ClientBuilder, Method, Request, RequestBuilder, Re use serde::Serialize; use tokio::sync::Mutex; -use crate::{deobfuscate::Deobfuscator, util}; +use crate::{ + cache::{Cache, DesktopClientData}, + deobfuscate::Deobfuscator, + util, +}; #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] pub enum ClientType { @@ -114,6 +118,7 @@ const IOS_DEVICE_MODEL: &str = "iPhone14,5"; pub struct RustyTube { pub locale: Arc, + cache: Cache, desktop_client: Arc, android_client: Arc, ios_client: Arc, @@ -128,19 +133,25 @@ pub struct Locale { impl RustyTube { #[must_use] pub fn new() -> Self { - Self::new_with_ua("en", "US") + Self::new_with_ua("en", "US", Some("rusty-tube.json".to_owned())) } #[must_use] - pub fn new_with_ua(lang: &str, country: &str) -> Self { + pub fn new_with_ua(lang: &str, country: &str, cache_file: Option) -> Self { let locale = Arc::new(Locale { lang: lang.to_owned(), country: country.to_owned(), }); + let cache = match cache_file.as_ref() { + Some(cache_file) => Cache::from_json_file(cache_file), + None => Cache::default(), + }; + Self { locale: locale.clone(), - desktop_client: Arc::new(DesktopClient::new(locale.clone())), + cache: cache.clone(), + desktop_client: Arc::new(DesktopClient::new(locale.clone(), cache)), android_client: Arc::new(AndroidClient::new(locale.clone())), ios_client: Arc::new(IosClient::new(locale)), } @@ -159,48 +170,29 @@ impl RustyTube { #[async_trait] pub trait YTClient { - // fn new(locale: Arc) -> Self; - async fn get_context(&self, localized: bool) -> ContextYT; async fn request_builder(&self, method: Method, url: &str) -> RequestBuilder; async fn exec_request(&self, request: Request) -> Result; async fn exec_request_text(&self, request: Request) -> Result; } +async fn exec_request(http: Client, request: Request) -> Result { + Ok(http.execute(request).await?.error_for_status()?) +} + +async fn exec_request_text(http: Client, request: Request) -> Result { + Ok(exec_request(http, request).await?.text().await?) +} + pub struct DesktopClient { locale: Arc, http: Client, - data: Mutex, + cache: Cache, consent_cookie_yes: String, consent_cookie_no: String, deobf: Deobfuscator, } -#[derive(Debug)] -struct DesktopClientData { - last_update: Option, - client_version: String, -} - -impl Default for DesktopClientData { - fn default() -> Self { - Self { - last_update: None, - client_version: DESKTOP_CLIENT_VERSION.to_owned(), - } - } -} - -impl DesktopClientData { - fn is_old(&self) -> bool { - self.last_update.is_none() - || Instant::now() - .duration_since(self.last_update.unwrap()) - .as_secs() - > 86400 - } -} - #[async_trait] impl YTClient for DesktopClient { async fn get_context(&self, localized: bool) -> ContextYT { @@ -253,7 +245,7 @@ impl YTClient for DesktopClient { } impl DesktopClient { - fn new(locale: Arc) -> Self { + fn new(locale: Arc, cache: Cache) -> Self { let mut rng = rand::thread_rng(); let http = ClientBuilder::new() @@ -263,12 +255,12 @@ impl DesktopClient { .build() .expect("unable to build the HTTP client"); - let deobf = Deobfuscator::new(http.clone()); + let deobf = Deobfuscator::new(http.clone(), cache.clone()); Self { locale, http, - data: Mutex::new(DesktopClientData::default()), + cache, consent_cookie_yes: format!( "{}={}{}", CONSENT_COOKIE, @@ -285,19 +277,21 @@ impl DesktopClient { } } - async fn extract_client_version_from_swjs(&self) -> Result> { - let swjs = self - .exec_request_text( - self.http - .get("https://www.youtube.com/sw.js") - .header(header::ORIGIN, "https://www.youtube.com") - .header(header::REFERER, "https://www.youtube.com") - .header(header::COOKIE, self.consent_cookie_yes.to_owned()) - .build() - .unwrap(), - ) - .await - .context("Failed to download sw.js")?; + async fn extract_client_version_from_swjs( + http: Client, + consent_cookie: &str, + ) -> Result { + let swjs = exec_request_text( + http.clone(), + http.get("https://www.youtube.com/sw.js") + .header(header::ORIGIN, "https://www.youtube.com") + .header(header::REFERER, "https://www.youtube.com") + .header(header::COOKIE, consent_cookie) + .build() + .unwrap(), + ) + .await + .context("Failed to download sw.js")?; static CLIENT_VERSION_PATTERNS: Lazy<[Regex; 3]> = Lazy::new(|| { [ @@ -307,41 +301,30 @@ impl DesktopClient { ] }); - Ok(util::get_cg_from_regexes( - CLIENT_VERSION_PATTERNS.iter(), - &swjs, - 1, - )) + util::get_cg_from_regexes(CLIENT_VERSION_PATTERNS.iter(), &swjs, 1) + .ok_or(anyhow!("Could not find desktop client version in sw.js")) } async fn get_client_version(&self) -> String { - let mut client_data = self.data.lock().await; + let http = self.http.clone(); + let consent_cookie = self.consent_cookie_yes.clone(); - if client_data.is_old() { - let client_version = self.extract_client_version_from_swjs().await; - let new_version = match client_version { - Ok(client_version) => match client_version { - Some(client_version) => { - debug!("Updated desktop client version to {}", client_version); - client_version - } - None => { - warn!("Could not find desktop client version in sw.js"); - DESKTOP_CLIENT_VERSION.to_owned() - } - }, - Err(e) => { - warn!("Could not extract desktop client version, Error: {}", e); - DESKTOP_CLIENT_VERSION.to_owned() - } - }; + let client_data = self + .cache + .get_desktop_client_data(async move { + let client_version = + Self::extract_client_version_from_swjs(http, &consent_cookie).await?; + Ok(DesktopClientData { client_version }) + }) + .await; - *client_data = DesktopClientData { - client_version: new_version, - last_update: Some(Instant::now()), + match client_data { + Ok(client_data) => client_data.client_version, + Err(e) => { + warn!("{}", e); + DESKTOP_CLIENT_VERSION.to_owned() } } - client_data.client_version.to_owned() } } diff --git a/src/client/player.rs b/src/client/player.rs index 989f04c..e995875 100644 --- a/src/client/player.rs +++ b/src/client/player.rs @@ -2,7 +2,7 @@ use anyhow::{anyhow, bail, Context, Result}; use reqwest::Method; use serde::Serialize; -use super::{response, ContextYT, ClientType, RustyTube, YTClient}; +use super::{response, ClientType, ContextYT, RustyTube, YTClient}; use crate::util; // REQUEST @@ -74,8 +74,7 @@ impl RustyTube { } }; - let resp = self - .desktop_client + let resp = client .request_builder(Method::POST, "player") .await .json(&request_body) @@ -97,7 +96,10 @@ mod tests { #[test(tokio::test)] async fn t_fetch_stream() { let rt = RustyTube::new(); - let stream = rt.fetch_player("ZeerrnuLi5E", ClientType::Desktop).await.unwrap(); + let stream = rt + .fetch_player("ZeerrnuLi5E", ClientType::Android) + .await + .unwrap(); dbg!(stream); } diff --git a/src/deobfuscate.rs b/src/deobfuscate.rs index 90be989..12aa74f 100644 --- a/src/deobfuscate.rs +++ b/src/deobfuscate.rs @@ -4,88 +4,63 @@ use log::debug; use once_cell::sync::Lazy; use reqwest::Client; use std::result::Result::Ok; -use std::time::Instant; -use tokio::sync::RwLock; +use crate::cache::{Cache, DeobfData}; use crate::util; pub struct Deobfuscator { http: Client, - cache: RwLock, -} - -#[derive(Default)] -struct JSCache { - last_update: Option, - js_url: String, - sig_fn: String, - nsig_fn: String, - sts: String, + cache: Cache, } impl Deobfuscator { #[must_use] - pub fn new(http: Client) -> Self { - Self { - http, - cache: RwLock::new(JSCache::default()), - } + pub fn new(http: Client, cache: Cache) -> Self { + Self { http, cache } } - async fn update(&self) -> Result<()> { - let mut cache = self.cache.write().await; + async fn get_deobf_data(&self) -> Result { + let http = self.http.clone(); - if cache.is_stale() { - let url = get_player_js_url(&self.http) - .await - .context("Failed to retrieve player.js URL")?; + self.cache + .get_deobf_data(async move { + let js_url = get_player_js_url(&http) + .await + .context("Failed to retrieve player.js URL")?; - let player_js = get_response(&self.http, &url) - .await - .context("Failed to download player.js")?; + let player_js = get_response(&http, &js_url) + .await + .context("Failed to download player.js")?; - debug!("Downloaded player.js from {}", url); + debug!("Downloaded player.js from {}", js_url); - let sig_fn = get_sig_fn(&player_js)?; - let nsig_fn = get_nsig_fn(&player_js)?; - let sts = get_sts(&player_js)?; + let sig_fn = get_sig_fn(&player_js)?; + let nsig_fn = get_nsig_fn(&player_js)?; + let sts = get_sts(&player_js)?; - *cache = JSCache { - last_update: Some(Instant::now()), - js_url: url.to_owned(), - sig_fn, - nsig_fn, - sts, - }; - } - Ok(()) + Ok(DeobfData { + js_url, + nsig_fn, + sig_fn, + sts, + }) + }) + .await } pub async fn deobfuscate_sig(&self, sig: &str) -> Result { - self.update().await?; - let cache = self.cache.read().await; - deobfuscate_sig(sig, &cache.sig_fn) + let deobf_data = self.get_deobf_data().await?; + deobfuscate_sig(sig, &deobf_data.sig_fn) } pub async fn deobfuscate_nsig(&self, nsig: &str) -> Result { - self.update().await?; - let cache = self.cache.read().await; - deobfuscate_nsig(nsig, &cache.nsig_fn) + let deobf_data = self.get_deobf_data().await?; + deobfuscate_nsig(nsig, &deobf_data.nsig_fn) } pub async fn get_sts(&self) -> Result { - self.update().await?; - let cache = self.cache.read().await; - Ok(cache.sts.to_owned()) - } -} - -impl JSCache { - fn is_stale(&self) -> bool { - match self.last_update { - Some(last_update) => Instant::now().duration_since(last_update).as_secs() > 3600, - None => true, - } + let deobf_data = self.get_deobf_data().await?; + Ok(deobf_data.sts) } } @@ -407,7 +382,8 @@ c[36](c[8],c[32]),c[20](c[25],c[10]),c[2](c[22],c[8]),c[32](c[20],c[16]),c[32](c #[test(tokio::test)] async fn t_update() { let client = Client::new(); - let deobf = Deobfuscator::new(client); + let cache = Cache::default(); + let deobf = Deobfuscator::new(client, cache); let deobf_sig = deobf.deobfuscate_sig("GOqGOqGOq0QJ8wRAIgaryQHfplJ9xJSKFywyaSMHuuwZYsoMTAvRvfm51qIGECIA5061zWeyfMPX9hEl_U6f9J0tr7GTJMKyPf5XNrJb5fb5i").await.unwrap(); println!("{}", deobf_sig); @@ -416,7 +392,8 @@ c[36](c[8],c[32]),c[20](c[25],c[10]),c[2](c[22],c[8]),c[32](c[20],c[16]),c[32](c #[test(tokio::test)] async fn t_parallel() { let client = Client::new(); - let deobf = Deobfuscator::new(client); + let cache = Cache::default(); + let deobf = Deobfuscator::new(client, cache); let deobf_arc = Arc::new(deobf); let (deobf_sig, deobf_nsig) = tokio::join!(