modules/
downloader.rs

1use std::{env, error::Error, path::Path, fs::File, io::prelude::*, io::BufWriter, time::Instant};
2use colored::*;
3use url::Url;
4use indicatif::ProgressBar;
5use async_trait::async_trait;
6use crate::utils::*;
7
8#[async_trait]
9pub trait Downloader {
10    async fn download(&self, url: &str, output_file: &Path) -> Result<(), Box<dyn Error>>;
11    async fn download_ts_file(&self, url: &str, output_file: &mut BufWriter<File>, pb: &mut ProgressBar) -> Result<(), Box<dyn Error>> {
12        let response = reqwest::get(url).await?;
13        let body = response.bytes().await?;
14
15        pb.inc(1);
16        output_file.write_all(&body)?;
17
18        Ok(())
19    }
20}
21
22pub struct SimpleDownloader;
23
24#[async_trait]
25impl Downloader for SimpleDownloader {
26    async fn download(&self, url: &str, output_file: &Path) -> Result<(), Box<dyn Error>> {
27        let mut response = reqwest::get(url).await?;
28
29        let length = response.content_length().unwrap_or(0);
30        let pb = create_progress_bar(length);
31    
32        let file = File::create(output_file)?;
33        let mut output_file = BufWriter::new(file);
34        while let Some(chunk) = response.chunk().await? {
35            output_file.write_all(&chunk)?;
36            pb.inc(chunk.len() as u64);
37        }
38        output_file.flush()?;
39    
40        pb.finish_and_clear();
41    
42        Ok(())
43    }
44}
45
46pub struct M3U8Downloader;
47
48#[async_trait]
49impl Downloader for M3U8Downloader {
50    async fn download(&self, url: &str, output_file: &Path) -> Result<(), Box<dyn Error>> {
51        let base_url = Url::parse(url)?;
52        let response = reqwest::get(base_url.join(url)?).await?;
53        let body = response.text().await?;
54        let mut lines = body.lines();
55
56        let mut ts_files = Vec::new();
57        while let Some(line) = lines.next() {
58            if line.starts_with("#EXT-X-STREAM-INF") {
59                lines.next();
60            } else if line.ends_with(".ts") {
61                ts_files.push(base_url.join(line)?);
62            }
63        }
64
65        let mut pb = create_progress_bar(ts_files.len() as u64);
66        let mut output_file = BufWriter::new(File::create(output_file)?);
67        for ts_file in ts_files {
68            self.download_ts_file(&ts_file.to_string(), &mut output_file, &mut pb).await?;
69        }
70        output_file.flush()?;
71
72        pb.finish_and_clear();
73
74        Ok(())
75    }
76}
77
78pub struct DownloadManager;
79
80impl DownloadManager {
81    pub fn get_downloader(type_: &str) -> Option<Box<dyn Downloader>> {
82        match type_ {
83            "twitch-clip" => Some(Box::new(SimpleDownloader)),
84            "twitch-video" => Some(Box::new(M3U8Downloader)),
85            "youtube-video" | "youtube-short" => Some(Box::new(SimpleDownloader)),
86            "tiktok-video" => Some(Box::new(SimpleDownloader)),
87            _ => None,
88        }
89    }
90
91    pub async fn download(&self, type_: &str, url: &str, title: &str) -> Result<(), Box<dyn Error>> {
92        println!("{} {}", format!("Downloading {}:", type_).blue(), title);
93
94        let url_type: Vec<&str> = type_.split('-').collect();
95        let title = remove_not_characters(title);
96        let output = format!("{}\\{}s\\{}.mp4", url_type[0], url_type[1], title);
97        let start = Instant::now();
98
99        if let Some(downloader) = Self::get_downloader(type_) {
100            downloader.download(url, Path::new(&output)).await?;
101        } else {
102            return Err("Invalid type".into());
103        }
104
105        println!("{} {}", format!("Downloaded {} in:", url_type[1]).blue(), get_elapsed_time(start));
106
107        let path = env::current_dir()?.join(output);
108        println!("{} {}", "Saved to:".blue(), format!("\x1B]8;;file://{}\x07{}\x1B]8;;\x07", path.display(), path.display()));
109
110        Ok(())
111    }
112}