diff --git a/Cargo.toml b/Cargo.toml index 990eead..f011e9f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,16 +16,18 @@ url = "2.5.8" reqwest = { version = "0.13.2", features = ["rustls", "stream", "socks", "cookies"], default-features = false } clap = { version = "4.6.0", features = ["derive"] } serde = { version = "1.0.228", features = ["derive"] } -serde_json = "1.0.149" +serde_json = { version = "1.0.149", features = ["raw_value"] } tokio = { version = "1.50.0", features = ["rt-multi-thread", "macros", "signal"] } tokio-util = { version = "0.7.18", features = ["io"] } futures = "0.3.32" warp = { version = "0.4.2", features = ["server"] } -chrono = "0.4.44" +chrono = { version = "0.4.44", features = ["serde"] } anyhow = "1.0.102" indexmap = { version = "2.13.0", features = ["serde"] } sha1 = "0.10.6" thiserror = "2.0.18" +har = "0.8.1" +base64 = "0.22.1" [dev-dependencies] insta = "1.46.3" diff --git a/NOTICE.md b/NOTICE.md index aefd566..41a8823 100644 --- a/NOTICE.md +++ b/NOTICE.md @@ -3,12 +3,18 @@ This software bundles the following third-party libraries: - anyhow version 1.0.102 (https://github.com/dtolnay/anyhow) Author: David Tolnay License: MIT OR Apache-2.0 + - base64 version 0.22.1 (https://github.com/marshallpierce/rust-base64) + Author: Marshall Pierce + License: MIT OR Apache-2.0 - chrono version 0.4.44 (https://github.com/chronotope/chrono) License: MIT OR Apache-2.0 - clap version 4.6.0 (https://github.com/clap-rs/clap) License: MIT OR Apache-2.0 - futures version 0.3.32 (https://github.com/rust-lang/futures-rs) License: MIT OR Apache-2.0 + - har version 0.8.1 (https://github.com/mandrean/har-rs) + Author: Sebastian Mandrean + License: MIT - indexmap version 2.13.0 (https://github.com/indexmap-rs/indexmap) License: Apache-2.0 OR MIT - m3u8-rs version 6.0.0 (https://github.com/rutgersc/m3u8-rs) diff --git a/src/import/har.rs b/src/import/har.rs new file mode 100644 index 0000000..d123e41 --- /dev/null +++ b/src/import/har.rs @@ -0,0 +1,133 @@ +use base64::prelude::*; +use chrono::{DateTime, Utc}; +pub use har::v1_3::*; +use serde::Deserialize; +use serde_json::value::RawValue; +use std::io::{Error, Result}; + +#[derive(Clone, Debug, Deserialize, Default)] +pub enum Version { + /// Version 1.2 of the HAR specification. + /// + /// Refer to the official + /// [specification](https://w3c.github.io/web-performance/specs/HAR/Overview.html) + /// for more information. + #[allow(non_camel_case_types)] + #[serde(rename = "1.2")] + #[default] + V1_2, + + // Version 1.3 of the HAR specification. + // + // Refer to the draft + // [specification](https://github.com/ahmadnassri/har-spec/blob/master/versions/1.3.md) + // for more information. + #[allow(non_camel_case_types)] + #[serde(rename = "1.3")] + V1_3, +} + +#[derive(Clone, Debug, Deserialize)] +pub struct Har { + pub log: Log, +} + +#[derive(Clone, Debug, Deserialize, Default)] +pub struct Log { + pub version: Version, + pub creator: Creator, + pub browser: Option, + pub pages: Option>, + pub entries: Vec, + pub comment: Option, +} + +#[derive(Clone, Debug, Deserialize, Default)] +pub struct Pages { + #[serde(rename = "startedDateTime")] + pub started_date_time: DateTime, + pub id: String, + pub title: String, + #[serde(rename = "pageTimings")] + pub page_timings: PageTimings, + pub comment: Option, +} + +#[derive(Clone, Debug, Deserialize, Default)] +pub struct Entries { + pub pageref: Option, + #[serde(rename = "startedDateTime")] + pub started_date_time: DateTime, + pub time: f64, + pub request: Request, + pub response: Response, + pub cache: Cache, + pub timings: Timings, + #[serde(rename = "serverIPAddress")] + pub server_ip_address: Option, + pub connection: Option, + pub comment: Option, +} + +#[derive(Clone, Debug, Deserialize, Default)] +pub struct Response { + pub status: i64, + #[serde(rename = "statusText")] + pub status_text: String, + #[serde(rename = "httpVersion")] + pub http_version: String, + pub cookies: Vec, + pub headers: Vec, + pub content: Content, + #[serde(rename = "redirectURL")] + pub redirect_url: Option, + #[serde(rename = "headersSize", default = "default_isize")] + pub headers_size: i64, + #[serde(rename = "bodySize", default = "default_isize")] + pub body_size: i64, + pub comment: Option, + #[serde(rename = "headersCompression")] + pub headers_compression: Option, +} + +#[derive(Clone, Debug, Deserialize, Default)] +pub struct Content { + #[serde(default = "default_isize")] + pub size: i64, + pub compression: Option, + #[serde(rename = "mimeType")] + pub mime_type: Option, + // CHANGED: Parse as raw value, to avoid keeping a ton of strings in memory. + // pub text: Option, + text: Option>, + pub encoding: Option, + pub comment: Option, +} + +impl Content { + pub(crate) fn text(&self) -> Result { + let Some(raw_text) = self.text.as_ref() else { + // Response has no content. + return Ok(String::new()); + }; + // Parse back into a string. + let text = serde_json::from_str::(raw_text.as_ref().get()) + .map_err(|_| Error::other("invalid content"))?; + Ok(text) + } + + pub(crate) fn as_bytes(&self) -> Result> { + let text = self.text()?; + // Decode as base64 (if needed). + match self.encoding.as_ref() { + Some(encoding) if encoding == "base64" => BASE64_STANDARD + .decode(text.as_bytes()) + .map_err(|_| Error::other("invalid base64")), + _ => Ok(text.into_bytes()), + } + } +} + +fn default_isize() -> i64 { + -1 +} diff --git a/src/import/mod.rs b/src/import/mod.rs new file mode 100644 index 0000000..b064985 --- /dev/null +++ b/src/import/mod.rs @@ -0,0 +1,43 @@ +mod har; + +use crate::record::{HarSource, RecordError, RecordOptions, Recorder}; +use crate::shared::url_file_extension; +use anyhow::anyhow; +pub(crate) use har::Har; +use std::io::Error; +use std::path::Path; +use tokio::fs; +use tokio_util::sync::CancellationToken; +use url::Url; + +pub async fn import_har( + har: Har, + dest: &Path, + options: RecordOptions, + token: CancellationToken, +) -> Result<(), RecordError> { + fs::create_dir_all(dest).await?; + + // Find the first .m3u8 request + let first_playlist_request = har + .log + .entries + .iter() + .find(|entry| { + let url = Url::parse(&entry.request.url).unwrap(); + matches!(url_file_extension(&url), Some("m3u8")) + }) + .ok_or_else(|| Error::other("no playlist found"))?; + let url = first_playlist_request + .request + .url + .parse::() + .map_err(|e| RecordError::Parse(anyhow!("Invalid URL: {e}")))?; + let time = first_playlist_request.started_date_time; + + // Create a source that reads from the HAR + let source = HarSource::new(har, time); + + let recorder = Recorder::new(source, dest, options).await?; + recorder.run(&url, token).await +} diff --git a/src/lib.rs b/src/lib.rs index dd90264..fd8ea08 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,3 +1,4 @@ +pub mod import; pub mod record; pub mod replay; pub mod shared; diff --git a/src/main.rs b/src/main.rs index c0026db..3308bf1 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,10 @@ +use std::fs::File; +use std::io::BufReader; use std::net::IpAddr; use std::path::PathBuf; use std::str::FromStr; -use clap::{Parser, Subcommand}; +use clap::{Args, Parser, Subcommand}; use reqwest::header::{HeaderMap, HeaderName, HeaderValue}; use tokio::spawn; use tokio_util::sync::CancellationToken; @@ -32,6 +34,48 @@ struct CliRequired { cli: Cli, } +#[derive(Args)] +struct RecordOrImportArgs { + /// The variant stream(s) to record. + #[arg(short = 'v', long, default_value = "first")] + variant: VariantSelect, + /// The audio renditions(s) to record. + #[arg(long, default_value = "default")] + audio: MediaSelect, + /// The video renditions(s) to record. + #[arg(long, default_value = "default")] + video: MediaSelect, + /// The subtitle renditions(s) to record. + #[arg(long, default_value = "default")] + subtitle: MediaSelect, + /// The maximum bandwidth of the variant stream to record. + /// + /// Cannot be used when --variant is set. + #[arg(short = 'b', long, conflicts_with = "variant")] + bandwidth: Option, + /// The start time of the first segment to record, in seconds. + /// + /// - If positive, the start time counts from the start of the first media playlist. + /// - If negative, the start time counts from the end of the first media playlist. + /// - If unset, the recording starts at the first segment of the first media playlist. + #[arg(long, allow_hyphen_values = true, verbatim_doc_comment)] + start: Option, + /// The end time of the first segment to record, in seconds. + /// + /// - If positive, the end time counts from the start of the first media playlist. + /// - If negative, the end time counts from the end of the first media playlist. + /// - If unset, the recording stops at the last segment of the last media playlist. + #[arg(long, allow_hyphen_values = true, verbatim_doc_comment)] + end: Option, + /// Whether to preserve the original file names of playlists and segments. + /// + /// This may not be compatible with all streams. For example, if the segment URLs + /// only differ by their query (e.g. `https://example.com/segment?num=1`), + /// then all segments would be written to the same `segment` file. + #[arg(long)] + keep_names: bool, +} + #[derive(Subcommand)] enum CliCommand { /// Record a HLS VOD or live stream. @@ -42,49 +86,13 @@ enum CliCommand { /// The directory path to store the recording of the HLS stream. #[arg(value_name = "PATH")] recording_path: PathBuf, - /// The variant stream(s) to record. - #[arg(short = 'v', long, default_value = "first")] - variant: VariantSelect, - /// The audio renditions(s) to record. - #[arg(long, default_value = "default")] - audio: MediaSelect, - /// The video renditions(s) to record. - #[arg(long, default_value = "default")] - video: MediaSelect, - /// The subtitle renditions(s) to record. - #[arg(long, default_value = "default")] - subtitle: MediaSelect, - /// The maximum bandwidth of the variant stream to record. - /// - /// Cannot be used when --variant is set. - #[arg(short = 'b', long, conflicts_with = "variant")] - bandwidth: Option, - /// The start time of the first segment to record, in seconds. - /// - /// - If positive, the start time counts from the start of the first media playlist. - /// - If negative, the start time counts from the end of the first media playlist. - /// - If unset, the recording starts at the first segment of the first media playlist. - #[arg(long, allow_hyphen_values = true, verbatim_doc_comment)] - start: Option, - /// The end time of the first segment to record, in seconds. - /// - /// - If positive, the end time counts from the start of the first media playlist. - /// - If negative, the end time counts from the end of the first media playlist. - /// - If unset, the recording stops at the last segment of the last media playlist. - #[arg(long, allow_hyphen_values = true, verbatim_doc_comment)] - end: Option, + #[command(flatten)] + args: RecordOrImportArgs, /// Custom HTTP header to send with all requests. /// /// Can be specified multiple times. Format: "Name: Value" #[arg(short = 'H', long = "header", value_name = "Name: Value", value_parser = parse_header)] headers: Vec<(HeaderName, HeaderValue)>, - /// Whether to preserve the original file names of playlists and segments. - /// - /// This may not be compatible with all streams. For example, if the segment URLs - /// only differ by their query (e.g. `https://example.com/segment?num=1`), - /// then all segments would be written to the same `segment` file. - #[arg(long)] - keep_names: bool, }, /// Replay a HLS VOD or live stream. Replay { @@ -98,16 +106,27 @@ enum CliCommand { #[arg(short = 'p', long, value_name = "PORT", default_value_t = 8080)] port: u16, }, + // Import an HTTP Archive (HAR) file as a recording. + Import { + /// The path of the HAR file. + #[arg(value_name = "HAR")] + har_path: PathBuf, + /// The directory path to store the recording of the HLS stream. + #[arg(value_name = "PATH")] + recording_path: PathBuf, + #[command(flatten)] + args: RecordOrImportArgs, + }, } #[tokio::main] -async fn main() { +async fn main() -> anyhow::Result<()> { let args = Cli::parse(); let Some(command) = args.command else { if args.license { println!(include_str!("../LICENSE.md")); println!(include_str!("../NOTICE.md")); - return; + return Ok(()); } else { // If --license is not set, then a subcommand is required. let Err(e) = CliRequired::try_parse() else { @@ -120,31 +139,16 @@ async fn main() { CliCommand::Record { manifest_url, recording_path, - variant, - audio, - video, - subtitle, - bandwidth, - start, - end, + args, headers, - keep_names, } => { - let variant_select = if let Some(bandwidth) = bandwidth { + let variant_select = if let Some(bandwidth) = args.bandwidth { VariantSelectOptions::Bandwidth(bandwidth) } else { - VariantSelectOptions::Named(variant) - }; - let options = RecordOptions { - start, - end, - variant_select, - audio, - video, - subtitle, - headers: headers.into_iter().collect::(), - keep_names, + VariantSelectOptions::Named(args.variant) }; + let headers = headers.into_iter().collect::(); + let options = args.into_record_options(variant_select, headers); let token = CancellationToken::new(); let record_task = { let token = token.clone(); @@ -182,7 +186,37 @@ async fn main() { } }; } + CliCommand::Import { + har_path, + recording_path, + args, + } => { + let variant_select = if let Some(bandwidth) = args.bandwidth { + VariantSelectOptions::Bandwidth(bandwidth) + } else { + VariantSelectOptions::Named(args.variant) + }; + let options = args.into_record_options(variant_select, HeaderMap::new()); + let token = CancellationToken::new(); + let har_file = File::open(har_path)?; + let har = serde_json::from_reader(BufReader::new(har_file))?; + let import_task = { + let token = token.clone(); + spawn(async move { + streamrr::import::import_har(har, &recording_path, options, token).await + }) + }; + match abort_on_ctrlc(import_task, token, RecordError::Cancelled).await { + Ok(()) => {} + Err(RecordError::Cancelled) => println!("Stopped importing."), + Err(e) => { + eprintln!("{e}"); + std::process::exit(1); + } + }; + } } + Ok(()) } fn parse_header(s: &str) -> Result<(HeaderName, HeaderValue), String> { @@ -195,3 +229,22 @@ fn parse_header(s: &str) -> Result<(HeaderName, HeaderValue), String> { HeaderValue::from_str(value.trim()).map_err(|e| e.to_string())?, )) } + +impl RecordOrImportArgs { + fn into_record_options( + self, + variant_select: VariantSelectOptions, + headers: HeaderMap, + ) -> RecordOptions { + RecordOptions { + start: self.start, + end: self.end, + variant_select, + audio: self.audio, + video: self.video, + subtitle: self.subtitle, + headers, + keep_names: self.keep_names, + } + } +} diff --git a/src/record/mod.rs b/src/record/mod.rs index cd14566..b7d6186 100644 --- a/src/record/mod.rs +++ b/src/record/mod.rs @@ -1,27 +1,29 @@ use anyhow::anyhow; use chrono::{DateTime, Utc}; -use futures::future::{BoxFuture, FutureExt, TryFutureExt}; +use futures::future::{BoxFuture, FutureExt}; use futures::stream::{StreamExt, TryStreamExt, iter}; use m3u8_rs::*; +use reqwest::Client; use reqwest::header::HeaderMap; -use reqwest::{Client, Response}; use std::io; use std::path::{Path, PathBuf}; +use std::pin::pin; use std::sync::Arc; -use std::time::{Duration, Instant}; +use std::time::Duration; use tokio::fs; use tokio::io::{AsyncSeekExt, AsyncWriteExt}; use tokio::sync::Mutex; use tokio::task::JoinSet; -use tokio::time::sleep_until; use tokio_util::io::StreamReader; use tokio_util::sync::CancellationToken; use url::Url; -use crate::shared::{ByteRange, MediaSelect, Recording, StripBom, VariantSelectOptions}; +use crate::shared::{ByteRange, MediaSelect, Recording, StripBom, Timed, VariantSelectOptions}; pub use rewrite::*; +pub use source::*; mod rewrite; +mod source; const MAX_CONCURRENT_DOWNLOADS: usize = 4; @@ -58,299 +60,312 @@ pub async fn record( options: RecordOptions, token: CancellationToken, ) -> Result<(), RecordError> { - fs::create_dir_all(dest).await?; - let recording_path = dest.join("recording.json"); - let recording = RecordingFile::new(&recording_path).await?; - let recording = Arc::new(Mutex::new(recording)); let client = Client::builder() .cookie_store(true) .default_headers(options.headers.clone()) .build() .map_err(|_| RecordError::Config("Error while building HTTP client"))?; - // Download initial playlist - let raw_playlist = token - .run_until_cancelled(download_playlist(&client, url)) - .await - .ok_or(RecordError::Cancelled)?? - .strip_bom(); - let initial_playlist = parse_playlist_res(raw_playlist.as_bytes()).map_err(|e| { - RecordError::Parse(anyhow!( - "Error while parsing playlist: {}", - e.map_input(|i| String::from_utf8_lossy(i)) - )) - })?; - match initial_playlist { - Playlist::MasterPlaylist(master_playlist) => { - // Master playlist - record_master_playlist( - &client, - url, - dest, - recording, - options, - master_playlist, - token, - ) - .await?; - } - Playlist::MediaPlaylist(media_playlist) => { - // Media playlist only - record_media_playlist( - &client, - url, - "", - Some(media_playlist), - dest, - recording, - options, - token, - ) - .await?; - } - } - Ok(()) + let source = HttpSource::new(client); + let recorder = Recorder::new(source, dest, options).await?; + recorder.run(url, token).await } -async fn record_master_playlist( - client: &Client, - url: &Url, - dest: &Path, +#[derive(Clone)] +pub struct Recorder { + source: S, + dest: PathBuf, recording: Arc>, options: RecordOptions, - mut master_playlist: MasterPlaylist, - token: CancellationToken, -) -> Result<(), RecordError> { - // Rewrite master playlist - let rewriter = Rewriter::new(url, dest, options.keep_names); - rewriter.rewrite_master_playlist(&mut master_playlist)?; - - // Select variant streams - master_playlist.variants = options - .variant_select - .filter_variants(&master_playlist.variants) - .to_vec(); - if master_playlist.variants.is_empty() { - return Err(RecordError::Config("No variant streams selected.")); - } - // Select renditions - let alternatives = &mut master_playlist.alternatives; - alternatives.retain(|media| { - // Must apply to at least one selected variant stream - master_playlist - .variants - .iter() - .any(|variant| media_applies_to_variant(media, variant)) - }); - let audio_renditions = alternatives - .iter() - .filter(|media| media.media_type == AlternativeMediaType::Audio) - .cloned() - .collect::>(); - let video_renditions = alternatives - .iter() - .filter(|media| media.media_type == AlternativeMediaType::Video) - .cloned() - .collect::>(); - let subtitle_renditions = alternatives - .iter() - .filter(|media| media.media_type == AlternativeMediaType::Subtitles) - .cloned() - .collect::>(); - let cc_renditions = alternatives - .iter() - .filter(|media| media.media_type == AlternativeMediaType::ClosedCaptions) - .cloned() - .collect::>(); - let other_renditions = alternatives - .iter() - .filter(|media| matches!(media.media_type, AlternativeMediaType::Other(_))) - .cloned() - .collect::>(); - - let audio_renditions = options.audio.filter_media(&audio_renditions).to_vec(); - let video_renditions = options.video.filter_media(&video_renditions).to_vec(); - let subtitle_renditions = options.subtitle.filter_media(&subtitle_renditions).to_vec(); - - master_playlist.alternatives = audio_renditions; - master_playlist.alternatives.extend(video_renditions); - master_playlist.alternatives.extend(subtitle_renditions); - master_playlist.alternatives.extend(cc_renditions); - master_playlist.alternatives.extend(other_renditions); - - let master_playlist = master_playlist; - - // Start recording selected variant streams and renditions - let mut join_set = JoinSet::new(); - for variant in &master_playlist.variants { - let Some(other_attributes) = &variant.other_attributes else { - continue; - }; - let Some(variant_url) = other_attributes.get(ORIGINAL_URI) else { - continue; - }; - let variant_url = Url::parse(variant_url.as_str()).unwrap(); - let variant_dir = Path::new(&variant.uri) - .parent() - .unwrap() - .to_string_lossy() - .to_string(); - let client = client.clone(); - let dest = PathBuf::from(dest); - let recording = recording.clone(); - let options = options.clone(); - let token = token.clone(); - join_set.spawn(async move { - record_media_playlist( - &client, - &variant_url, - &variant_dir, - None, - &dest, - recording, - options, - token, - ) - .await - }); +} + +impl Recorder { + pub async fn new(source: S, dest: &Path, options: RecordOptions) -> Result { + fs::create_dir_all(dest).await?; + let recording_path = dest.join("recording.json"); + let recording = RecordingFile::new(&recording_path).await?; + let recording = Arc::new(Mutex::new(recording)); + Ok(Self { + source, + dest: PathBuf::from(dest), + recording, + options, + }) } - for media in &master_playlist.alternatives { - let Some(media_uri) = &media.uri else { - continue; - }; - let Some(other_attributes) = &media.other_attributes else { - continue; - }; - let Some(media_url) = other_attributes.get(ORIGINAL_URI) else { - continue; - }; - let media_url = Url::parse(media_url.as_str()).unwrap(); - let media_dir = Path::new(media_uri) - .parent() - .unwrap() - .to_string_lossy() - .to_string(); - let client = client.clone(); - let dest = PathBuf::from(dest); - let recording = recording.clone(); - let options = options.clone(); - let token = token.clone(); - join_set.spawn(async move { - record_media_playlist( - &client, &media_url, &media_dir, None, &dest, recording, options, token, - ) + + pub async fn run(mut self, url: &Url, token: CancellationToken) -> Result<(), RecordError> { + // Download initial playlist + let Timed { + value: initial_playlist, + time: playlist_time, + } = token + .run_until_cancelled(download_playlist(&self.source, url)) .await - }); - } - // Write updated master playlist - let master_name = "index.m3u8"; - write_master_playlist(&dest.join(master_name), &master_playlist).await?; - recording - .lock() - .await - .add_and_save(Utc::now(), master_name, master_name.to_string()) - .await?; - // Wait for all tasks to complete - while let Some(res) = join_set.join_next().await { - res.map_err(|_| RecordError::Cancelled)??; + .ok_or(RecordError::Cancelled)?? + .and_then(|raw_playlist| { + let raw_playlist = raw_playlist.strip_bom(); + parse_playlist_res(raw_playlist.as_bytes()).map_err(|e| { + RecordError::Parse(anyhow!( + "Error while parsing playlist: {}", + e.map_input(|i| String::from_utf8_lossy(i)) + )) + }) + })?; + match initial_playlist { + Playlist::MasterPlaylist(master_playlist) => { + // Master playlist + self.record_master_playlist(url, master_playlist, token) + .await?; + } + Playlist::MediaPlaylist(media_playlist) => { + // Media playlist only + self.record_media_playlist( + url, + "", + Some(Timed { + value: media_playlist, + time: playlist_time, + }), + token, + ) + .await?; + } + } + Ok(()) } - Ok(()) -} -#[allow(clippy::too_many_arguments)] -async fn record_media_playlist( - client: &Client, - url: &Url, - dir: &str, - mut initial_playlist: Option, - dest: &Path, - recording: Arc>, - mut options: RecordOptions, - token: CancellationToken, -) -> Result<(), RecordError> { - let dest_dir = dest.join(dir); - fs::create_dir_all(&dest_dir).await?; - let mut rewriter = Rewriter::new(url, dir.as_ref(), options.keep_names); - let name_in_recording = rewriter.playlist_path(); - let mut previous_playlist = None; - let mut lowest_media_sequence = 0; - let mut highest_media_sequence = None; - loop { - // Download and rewrite playlist - let mut media_playlist = if let Some(playlist) = initial_playlist.take() { - playlist - } else { - let raw_playlist = token - .run_until_cancelled(download_playlist(client, url)) - .await - .ok_or(RecordError::Cancelled)?? - .strip_bom(); - parse_media_playlist_res(raw_playlist.as_bytes()).map_err(|e| { - RecordError::Parse(anyhow!( - "Error while parsing media playlist: {}", - e.map_input(|i| String::from_utf8_lossy(i)) - )) - })? - }; - let now = Instant::now(); - let playlist_time = Utc::now(); - let file_name = if previous_playlist.is_none() && media_playlist.end_list { - // Playlist is a VOD. No need for a timestamp, since we won't ever refresh it. - rewriter.playlist_path() - } else { - // Playlist is live, or was live and has now ended - rewriter.playlist_path_with_timestamp(&playlist_time) - }; - // Clip to start and end time (if given) - if let Some(start) = options.start.take() - && let Some(start_index) = find_segment_index_by_offset(&media_playlist.segments, start) - { - lowest_media_sequence = media_playlist.media_sequence + (start_index as u64) + async fn record_master_playlist( + &self, + url: &Url, + mut master_playlist: MasterPlaylist, + token: CancellationToken, + ) -> Result<(), RecordError> { + // Rewrite master playlist + let rewriter = Rewriter::new(url, &self.dest, self.options.keep_names); + rewriter.rewrite_master_playlist(&mut master_playlist)?; + + // Select variant streams and renditions + self.select_variants(&mut master_playlist); + if master_playlist.variants.is_empty() { + return Err(RecordError::Config("No variant streams selected.")); } - if let Some(end) = options.end.take() - && let Some(end_index) = find_segment_index_by_offset(&media_playlist.segments, end) - { - highest_media_sequence = Some(media_playlist.media_sequence + (end_index as u64)) + self.select_renditions(&mut master_playlist); + let master_playlist = master_playlist; + + // Start recording selected variant streams and renditions + let mut join_set = JoinSet::new(); + for variant in &master_playlist.variants { + let Some(other_attributes) = &variant.other_attributes else { + continue; + }; + let Some(variant_url) = other_attributes.get(ORIGINAL_URI) else { + continue; + }; + let variant_url = Url::parse(variant_url.as_str()).unwrap(); + let variant_dir = Path::new(&variant.uri) + .parent() + .unwrap() + .to_string_lossy() + .to_string(); + let mut recorder = self.clone(); + let token = token.clone(); + join_set.spawn(async move { + recorder + .record_media_playlist(&variant_url, &variant_dir, None, token) + .await + }); } - remove_segments_from_start(&mut media_playlist, lowest_media_sequence); - if let Some(highest_media_sequence) = highest_media_sequence { - remove_segments_from_end(&mut media_playlist, highest_media_sequence); + for media in &master_playlist.alternatives { + let Some(media_uri) = &media.uri else { + continue; + }; + let Some(other_attributes) = &media.other_attributes else { + continue; + }; + let Some(media_url) = other_attributes.get(ORIGINAL_URI) else { + continue; + }; + let media_url = Url::parse(media_url.as_str()).unwrap(); + let media_dir = Path::new(media_uri) + .parent() + .unwrap() + .to_string_lossy() + .to_string(); + let mut recorder = self.clone(); + let token = token.clone(); + join_set.spawn(async move { + recorder + .record_media_playlist(&media_url, &media_dir, None, token) + .await + }); } - rewriter.rewrite_media_playlist(&mut media_playlist)?; - write_media_playlist(&dest.join(&file_name), &media_playlist).await?; - // Update recording - recording + // Write updated master playlist + let master_name = "index.m3u8"; + write_master_playlist(&self.dest.join(master_name), &master_playlist).await?; + self.recording .lock() .await - .add_and_save(playlist_time, &name_in_recording, file_name.to_string()) + .add_and_save(Utc::now(), master_name, master_name.to_string()) .await?; - // Download segments - download_segments( - client, - &media_playlist.segments, - &dest_dir, - MAX_CONCURRENT_DOWNLOADS, - token.clone(), - ) - .await?; - // Refresh playlist - if media_playlist.end_list { - break; + // Wait for all tasks to complete + while let Some(res) = join_set.join_next().await { + res.map_err(|_| RecordError::Cancelled)??; } - let next_refresh_time = now + Duration::from_secs(media_playlist.target_duration); - token - .run_until_cancelled(sleep_until(next_refresh_time.into())) - .await - .ok_or(RecordError::Cancelled)?; - previous_playlist = Some(media_playlist); + Ok(()) + } + + fn select_variants(&self, master_playlist: &mut MasterPlaylist) { + master_playlist.variants = self + .options + .variant_select + .filter_variants(&master_playlist.variants) + .to_vec(); + } + + fn select_renditions(&self, master_playlist: &mut MasterPlaylist) { + let alternatives = &mut master_playlist.alternatives; + alternatives.retain(|media| { + // Must apply to at least one selected variant stream + master_playlist + .variants + .iter() + .any(|variant| media_applies_to_variant(media, variant)) + }); + let audio_renditions = alternatives + .iter() + .filter(|media| media.media_type == AlternativeMediaType::Audio) + .cloned() + .collect::>(); + let video_renditions = alternatives + .iter() + .filter(|media| media.media_type == AlternativeMediaType::Video) + .cloned() + .collect::>(); + let subtitle_renditions = alternatives + .iter() + .filter(|media| media.media_type == AlternativeMediaType::Subtitles) + .cloned() + .collect::>(); + let cc_renditions = alternatives + .iter() + .filter(|media| media.media_type == AlternativeMediaType::ClosedCaptions) + .cloned() + .collect::>(); + let other_renditions = alternatives + .iter() + .filter(|media| matches!(media.media_type, AlternativeMediaType::Other(_))) + .cloned() + .collect::>(); + + let audio_renditions = self.options.audio.filter_media(&audio_renditions).to_vec(); + let video_renditions = self.options.video.filter_media(&video_renditions).to_vec(); + let subtitle_renditions = self + .options + .subtitle + .filter_media(&subtitle_renditions) + .to_vec(); + + master_playlist.alternatives = audio_renditions; + master_playlist.alternatives.extend(video_renditions); + master_playlist.alternatives.extend(subtitle_renditions); + master_playlist.alternatives.extend(cc_renditions); + master_playlist.alternatives.extend(other_renditions); + } + + async fn record_media_playlist( + &mut self, + url: &Url, + dir: &str, + mut initial_playlist: Option>, + token: CancellationToken, + ) -> Result<(), RecordError> { + let dest_dir = self.dest.join(dir); + fs::create_dir_all(&dest_dir).await?; + let mut rewriter = Rewriter::new(url, dir.as_ref(), self.options.keep_names); + let name_in_recording = rewriter.playlist_path(); + let mut previous_playlist = None; + let mut lowest_media_sequence = 0; + let mut highest_media_sequence = None; + loop { + // Download and rewrite playlist + let Timed { + value: mut media_playlist, + time: playlist_time, + } = if let Some(playlist) = initial_playlist.take() { + playlist + } else { + token + .run_until_cancelled(download_playlist(&self.source, url)) + .await + .ok_or(RecordError::Cancelled)?? + .and_then(|raw_playlist| { + let raw_playlist = raw_playlist.strip_bom(); + parse_media_playlist_res(raw_playlist.as_bytes()).map_err(|e| { + RecordError::Parse(anyhow!( + "Error while parsing media playlist: {}", + e.map_input(|i| String::from_utf8_lossy(i)) + )) + }) + })? + }; + let file_name = if previous_playlist.is_none() && media_playlist.end_list { + // Playlist is a VOD. No need for a timestamp, since we won't ever refresh it. + rewriter.playlist_path() + } else { + // Playlist is live, or was live and has now ended + rewriter.playlist_path_with_timestamp(&playlist_time) + }; + // Clip to start and end time (if given) + if let Some(start) = self.options.start.take() + && let Some(start_index) = + find_segment_index_by_offset(&media_playlist.segments, start) + { + lowest_media_sequence = media_playlist.media_sequence + (start_index as u64) + } + if let Some(end) = self.options.end.take() + && let Some(end_index) = find_segment_index_by_offset(&media_playlist.segments, end) + { + highest_media_sequence = Some(media_playlist.media_sequence + (end_index as u64)) + } + remove_segments_from_start(&mut media_playlist, lowest_media_sequence); + if let Some(highest_media_sequence) = highest_media_sequence { + remove_segments_from_end(&mut media_playlist, highest_media_sequence); + } + rewriter.rewrite_media_playlist(&mut media_playlist)?; + write_media_playlist(&self.dest.join(&file_name), &media_playlist).await?; + // Update recording + self.recording + .lock() + .await + .add_and_save(playlist_time, &name_in_recording, file_name.to_string()) + .await?; + // Download segments + download_segments( + &self.source, + &media_playlist.segments, + &dest_dir, + MAX_CONCURRENT_DOWNLOADS, + token.clone(), + ) + .await?; + // Refresh playlist + if media_playlist.end_list { + break; + } + let next_refresh_time = + playlist_time + Duration::from_secs(media_playlist.target_duration); + token + .run_until_cancelled(self.source.advance_to_time(next_refresh_time)) + .await + .ok_or(RecordError::Cancelled)?; + previous_playlist = Some(media_playlist); + } + Ok(()) } - Ok(()) } -async fn download_playlist(client: &Client, url: &Url) -> Result { - client - .get(url.clone()) - .send() - .and_then(Response::text) +async fn download_playlist(source: &S, url: &Url) -> Result, RecordError> { + source + .request_string(url, None) .await .map_err(|e| RecordError::Io(io::Error::other(e))) } @@ -401,8 +416,8 @@ fn find_segment_index_by_offset(segments: &[MediaSegment], offset: f32) -> Optio } } -async fn download_segments( - client: &Client, +async fn download_segments( + source: &S, media_segments: &[MediaSegment], dir: &Path, max_concurrent_downloads: usize, @@ -410,7 +425,7 @@ async fn download_segments( ) -> Result<(), RecordError> { let segment_tasks = media_segments .iter() - .flat_map(|segment| make_segment_download_tasks(client, dir, segment, token.clone())); + .flat_map(|segment| make_segment_download_tasks(source, dir, segment, token.clone())); iter(segment_tasks) .boxed() // https://github.com/rust-lang/rust/issues/104382 .buffered(max_concurrent_downloads) @@ -418,25 +433,25 @@ async fn download_segments( .await } -fn make_segment_download_tasks<'a>( - client: &'a Client, +fn make_segment_download_tasks<'a, S: Source>( + source: &'a S, dir: &'a Path, segment: &'a MediaSegment, token: CancellationToken, ) -> Vec>> { let mut tasks = Vec::with_capacity(3); - tasks.push(download_segment(client, segment, dir, token.clone()).boxed()); + tasks.push(download_segment(source, segment, dir, token.clone()).boxed()); if let Some(key) = segment.key.as_ref() { - tasks.push(download_key(client, key, segment, dir, token.clone()).boxed()); + tasks.push(download_key(source, key, segment, dir, token.clone()).boxed()); } if let Some(map) = segment.map.as_ref() { - tasks.push(download_map(client, map, segment, dir, token).boxed()); + tasks.push(download_map(source, map, segment, dir, token).boxed()); } tasks } -async fn download_segment( - client: &Client, +async fn download_segment( + source: &S, media_segment: &MediaSegment, dir: &Path, token: CancellationToken, @@ -460,7 +475,7 @@ async fn download_segment( })?; let segment_file = &media_segment.uri; download_file( - client, + source, segment_url, segment_byte_range, segment_file, @@ -470,8 +485,8 @@ async fn download_segment( .await } -async fn download_key( - client: &Client, +async fn download_key( + source: &S, key: &Key, media_segment: &MediaSegment, dir: &Path, @@ -487,7 +502,7 @@ async fn download_key( let key_uri = original_key_tag.rest.as_ref().unwrap(); let key_file = key.uri.as_ref().unwrap(); download_file( - client, + source, key_uri.as_str(), None, key_file.as_str(), @@ -497,8 +512,8 @@ async fn download_key( .await } -async fn download_map( - client: &Client, +async fn download_map( + source: &S, map: &Map, media_segment: &MediaSegment, dir: &Path, @@ -522,7 +537,7 @@ async fn download_map( })?; let map_file = &map.uri; download_file( - client, + source, map_uri.as_str(), map_byte_range, map_file, @@ -532,8 +547,8 @@ async fn download_map( .await } -async fn download_file( - client: &Client, +async fn download_file( + source: &S, url: &str, byte_range: Option, file_name: &str, @@ -551,26 +566,15 @@ async fn download_file( Err(err) if err.kind() == io::ErrorKind::AlreadyExists => return Ok(()), Err(err) => return Err(err.into()), }; - let range_header = byte_range.map(|byte_range| { - let start = byte_range.offset; - let end = start + byte_range.length - 1; // end byte for a range request is inclusive! - format!("bytes={start}-{end}") - }); - println!( - "Download: {url} {}", - range_header.as_ref().unwrap_or(&String::new()) - ); - let mut request = client.get(url); - if let Some(range_header) = range_header { - request = request.header(reqwest::header::RANGE, range_header); - } - let response = token - .run_until_cancelled(request.send()) + let url = url + .parse::() + .map_err(|e| RecordError::Parse(anyhow!("Invalid URL: {e}")))?; + let response_stream = token + .run_until_cancelled(source.request_stream(&url, byte_range)) .await .ok_or(RecordError::Cancelled)? - .map_err(|e| RecordError::Io(io::Error::other(e)))?; - let response_stream = response.bytes_stream().map_err(io::Error::other); - let mut response_stream = StreamReader::new(response_stream); + .map_err(io::Error::other); + let mut response_stream = pin!(StreamReader::new(response_stream)); tokio::io::copy_buf(&mut response_stream, &mut file).await?; Ok(()) } @@ -634,34 +638,34 @@ mod tests { fn require_async_fn_to_be_send() { let url = Url::parse("https://a.com/").unwrap(); let path = Path::new(""); - let client = Client::new(); + let source = HttpSource::new(Client::new()); let token = CancellationToken::new(); fn require_send(_t: T) {} - require_send(download_playlist(&client, &url)); + require_send(download_playlist(&source, &url)); require_send(write_master_playlist(path, &MasterPlaylist::default())); require_send(write_media_playlist(path, &MediaPlaylist::default())); - require_send(download_segments(&client, &[], path, 0, token.clone())); + require_send(download_segments(&source, &[], path, 0, token.clone())); require_send(download_segment( - &client, + &source, &MediaSegment::empty(), path, token.clone(), )); require_send(download_key( - &client, + &source, &Key::default(), &MediaSegment::empty(), path, token.clone(), )); require_send(download_map( - &client, + &source, &Map::default(), &MediaSegment::empty(), path, token.clone(), )); - require_send(download_file(&client, "", None, "", path, token.clone())); + require_send(download_file(&source, "", None, "", path, token.clone())); } } diff --git a/src/record/source.rs b/src/record/source.rs new file mode 100644 index 0000000..89d7cd3 --- /dev/null +++ b/src/record/source.rs @@ -0,0 +1,55 @@ +use crate::shared::ByteRange; +use chrono::{DateTime, Utc}; +use futures::{Stream, future::ready, stream::once}; +use tokio_util::bytes::Bytes; +use url::Url; + +pub mod har; +pub mod http; + +use crate::shared::Timed; +pub use har::HarSource; +pub use http::HttpSource; + +pub trait Source: Clone + Send + Sync { + type Error: Into> + Send; + + /// Advance the current time for subsequent requests. + #[allow(unused_variables)] + fn advance_to_time(&mut self, time: DateTime) -> impl Future + Send { + ready(()) + } + + /// Request the resource at the given URL + /// as it was at the time set by `advance_to_time`. + fn request( + &self, + url: &Url, + byte_range: Option, + ) -> impl Future, Self::Error>> + Send; + + /// Request the resource at the given URL as a string, + /// as it was at the time set by `advance_to_time`. + fn request_string( + &self, + url: &Url, + byte_range: Option, + ) -> impl Future, Self::Error>> + Send { + async move { + let bytes = self.request(url, byte_range).await?; + Ok(bytes.map(|bytes| String::from_utf8_lossy(&bytes).to_string())) + } + } + + /// Request the resource at the given URL as a stream. + fn request_stream( + &self, + url: &Url, + byte_range: Option, + ) -> impl Future> + Send> + Send { + async move { + let bytes = self.request(url, byte_range).await.map(|bytes| bytes.value); + once(ready(bytes)) + } + } +} diff --git a/src/record/source/har.rs b/src/record/source/har.rs new file mode 100644 index 0000000..26e3df5 --- /dev/null +++ b/src/record/source/har.rs @@ -0,0 +1,54 @@ +use crate::import::Har; +use crate::record::Source; +use crate::shared::{ByteRange, Timed}; +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use tokio_util::bytes::Bytes; +use url::Url; + +#[derive(Clone)] +pub struct HarSource { + har: Har, + time: DateTime, +} + +impl HarSource { + pub fn new(har: Har, time: DateTime) -> Self { + Self { har, time } + } +} + +impl Source for HarSource { + type Error = anyhow::Error; + + async fn advance_to_time(&mut self, time: DateTime) { + self.time = time; + } + + async fn request( + &self, + url: &Url, + byte_range: Option, + ) -> Result, Self::Error> { + let entries = &self.har.log.entries; + let entry = entries + .iter() + .filter(|entry| entry.started_date_time >= self.time) + .filter(|entry| entry.request.url == url.as_str()) + .min_by_key(|entry| entry.started_date_time) + .ok_or_else(|| anyhow!("Not found: {url}"))?; + let mut content = entry + .response + .content + .as_bytes() + .map_err(|_| anyhow!("Invalid data for: {url}"))?; + if let Some(byte_range) = byte_range { + content.drain(0..(byte_range.offset as usize)); + content.truncate(byte_range.length as usize); + } + Ok(Timed { + value: Bytes::from(content), + time: entry.started_date_time, + }) + } +} diff --git a/src/record/source/http.rs b/src/record/source/http.rs new file mode 100644 index 0000000..0034941 --- /dev/null +++ b/src/record/source/http.rs @@ -0,0 +1,84 @@ +use super::Source; +use crate::shared::{ByteRange, Timed}; +use chrono::{DateTime, Utc}; +use futures::{Stream, TryStreamExt, future::ready, stream::once}; +use reqwest::{Client, RequestBuilder, Response}; +use tokio::time::sleep; +use tokio_util::bytes::Bytes; +use url::Url; + +#[derive(Clone)] +pub struct HttpSource { + client: Client, +} + +impl HttpSource { + pub fn new(client: Client) -> Self { + Self { client } + } + + fn build_request(&self, url: &Url, byte_range: Option) -> RequestBuilder { + let range_header = byte_range.map(|byte_range| { + let start = byte_range.offset; + let end = start + byte_range.length - 1; // end byte for a range request is inclusive! + format!("bytes={start}-{end}") + }); + + println!( + "Download: {url} {}", + range_header.as_ref().unwrap_or(&String::new()) + ); + + let mut request = self.client.get(url.clone()); + if let Some(range) = range_header { + request = request.header(reqwest::header::RANGE, range); + } + request + } +} + +impl Source for HttpSource { + type Error = reqwest::Error; + + async fn advance_to_time(&mut self, time: DateTime) { + let now = Utc::now(); + let duration = (now - time).to_std().unwrap_or_default(); + sleep(duration).await + } + + async fn request( + &self, + url: &Url, + byte_range: Option, + ) -> Result, Self::Error> { + let request = self.build_request(url, byte_range); + let response = request.send().await?; + let time = Utc::now(); + let bytes = response.bytes().await?; + Ok(Timed { value: bytes, time }) + } + + async fn request_string( + &self, + url: &Url, + byte_range: Option, + ) -> Result, Self::Error> { + let request = self.build_request(url, byte_range); + let response = request.send().await?; + let time = Utc::now(); + let text = response.text().await?; + Ok(Timed { value: text, time }) + } + + async fn request_stream( + &self, + url: &Url, + byte_range: Option, + ) -> impl Stream> { + let request = self.build_request(url, byte_range); + let response = request.send().await; + once(ready(response)) + .map_ok(Response::bytes_stream) + .try_flatten() + } +} diff --git a/src/shared/mod.rs b/src/shared/mod.rs index 6d8f5bb..87c8b9b 100644 --- a/src/shared/mod.rs +++ b/src/shared/mod.rs @@ -3,6 +3,7 @@ pub use byte_range::*; pub use ctrlc::*; pub use hexstring::*; pub use recording::*; +pub use timed::*; pub(crate) use url::*; mod bom; @@ -10,4 +11,5 @@ mod byte_range; mod ctrlc; mod hexstring; mod recording; +mod timed; mod url; diff --git a/src/shared/timed.rs b/src/shared/timed.rs new file mode 100644 index 0000000..6b4c201 --- /dev/null +++ b/src/shared/timed.rs @@ -0,0 +1,23 @@ +use chrono::{DateTime, Utc}; + +/// A value with a timestamp. +pub struct Timed { + pub value: T, + pub time: DateTime, +} + +impl Timed { + pub fn map(self, f: impl FnOnce(T) -> U) -> Timed { + Timed { + value: f(self.value), + time: self.time, + } + } + + pub fn and_then(self, f: impl FnOnce(T) -> Result) -> Result, E> { + Ok(Timed { + value: f(self.value)?, + time: self.time, + }) + } +}