use anyhow::Result;
use indicatif::{ProgressBar, ProgressStyle};
use nom::Finish;
use std::{
fs,
io::{self, copy, Read},
path::Path,
};
use url::Url;
use crate::agent::{AgentBase, YaydlAgent};
struct DownloadProgress<'a, R> {
inner: R,
progress_bar: &'a ProgressBar,
}
impl<R: Read> Read for DownloadProgress<'_, R> {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
self.inner.read(buf).map(|n| {
self.progress_bar.inc(n as u64);
n
})
}
}
pub fn download_from_playlist(url: &str, filename: &str, verbose: bool) -> Result<()> {
if verbose {
println!("{}", "Found a playlist. Fetching ...");
}
let mut url = Url::parse(url)?;
let agent = YaydlAgent::init(url.clone());
let request = agent.get(url.as_str()).header("Referer", url.as_str());
let playlist_text = request
.call()
.expect("Could not go to the playlist url")
.body_mut()
.read_to_string()
.expect("Could not read the playlist source");
if verbose {
println!("{}", "Parsing ...");
}
let playlist = m3u8_rs::parse_media_playlist(&playlist_text.as_bytes())
.finish()
.unwrap();
let file = Path::new(&filename);
let mut dest = fs::OpenOptions::new()
.create(true)
.append(true)
.open(&file)?;
let total_cnt = playlist.1.segments.len() as u64;
let pb = ProgressBar::new(total_cnt);
pb.set_style(
ProgressStyle::with_template(
"{spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {percent}%",
)
.unwrap()
.progress_chars("#>-"),
);
for segment in &playlist.1.segments {
url.path_segments_mut().unwrap().pop().push(&segment.uri);
let mut request = agent
.get(url.as_str())
.header("Referer", url.as_str())
.call()?;
let mut source = request.body_mut().as_reader();
let _ = copy(&mut source, &mut dest)?;
pb.inc(1);
}
pb.finish_and_clear();
Ok(())
}
pub fn download(url: &str, filename: &str) -> Result<()> {
let url = Url::parse(url)?;
let agent = YaydlAgent::init(url.clone());
let mut resp = agent
.get(url.as_str())
.header("Referer", url.as_str())
.call()?;
let total_size = resp
.body()
.content_length()
.expect("Failed to read the segment size. Aborting.");
let pb = ProgressBar::new(total_size);
pb.set_style(
ProgressStyle::with_template(
"{spinner:.green} [{elapsed_precise}] [{bar:40.green/blue}] {percent}%",
)
.unwrap()
.progress_chars("#>-"),
);
let file = Path::new(filename);
if file.exists() {
let size = file.metadata()?.len() - 1;
resp = agent
.get(url.as_str())
.header("Referer", url.as_str())
.header("Range", &format!("bytes={}-", size))
.call()?;
pb.inc(size);
}
let mut source = DownloadProgress {
progress_bar: &pb,
inner: resp.body_mut().as_reader(),
};
let mut dest = fs::OpenOptions::new()
.create(true)
.append(true)
.open(&file)?;
let _ = copy(&mut source, &mut dest)?;
pb.finish_and_clear();
Ok(())
}