working v1
This commit is contained in:
Oliver Atkinson
256
src/main.rs
Normal file
256
src/main.rs
Normal file
@@ -0,0 +1,256 @@
|
||||
#![deny(clippy::unwrap_used)]
|
||||
#![deny(clippy::todo)]
|
||||
#![deny(clippy::unimplemented)]
|
||||
#![feature(iter_collect_into)]
|
||||
|
||||
use std::{collections::HashMap, process::Command, sync::Arc};
|
||||
|
||||
use poem::{
|
||||
endpoint::StaticFilesEndpoint, get, handler, http::StatusCode, listener::TcpListener, middleware::AddData, post, web::{Data, Form, Path, StaticFileRequest}, Body, EndpointExt, FromRequest, IntoResponse, Request, Route, Server
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use tokio::sync::RwLock;
|
||||
use tracing::{debug, error, info, instrument, trace, Level};
|
||||
use tracing_subscriber::EnvFilter;
|
||||
use url::Url;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), std::io::Error> {
|
||||
|
||||
#[cfg(debug_assertions)]
|
||||
let filter = EnvFilter::builder()
|
||||
.parse("server=trace,poem=debug,tokio=warn")
|
||||
.expect("Could not create env filter.")
|
||||
;
|
||||
|
||||
#[cfg(debug_assertions)]
|
||||
tracing_subscriber::fmt::fmt()
|
||||
.with_max_level(Level::TRACE)
|
||||
.with_target(true)
|
||||
.with_env_filter(filter)
|
||||
.with_thread_ids(false)
|
||||
.with_file(false)
|
||||
.without_time()
|
||||
.init();
|
||||
|
||||
if !(check_cmd("ffmpeg", "-version") && check_cmd("yt-dlp", "--version")) {
|
||||
error!("One or more commands errored out, continuing is pointless.");
|
||||
return Ok(());
|
||||
}
|
||||
trace!("Loading up server...");
|
||||
run_server().await
|
||||
}
|
||||
|
||||
fn check_cmd(cmd: &str, ver_arg: &str) -> bool {
|
||||
match Command::new(cmd).arg(ver_arg).output()
|
||||
{
|
||||
Ok(e) => {
|
||||
let stderr = e.stderr.iter().map(|f| *f as char).collect::<String>();
|
||||
let stdout = e.stdout.iter().map(|f| *f as char).collect::<String>();
|
||||
|
||||
if !e.status.success() {
|
||||
error!("Failed with to run {}!", cmd);
|
||||
error!("{}", stderr);
|
||||
false
|
||||
} else {
|
||||
info!("{} version: {}", cmd, stdout);
|
||||
true
|
||||
}
|
||||
},
|
||||
Err(_) => {
|
||||
error!("Couldn't run commands at all!");
|
||||
false
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct AppData {
|
||||
mapping: RwLock<HashMap<String, std::path::PathBuf>>,
|
||||
}
|
||||
|
||||
pub static FILESYSTEM_ROOT: &str = "./www/";
|
||||
#[instrument(level="trace", name="web" skip_all)]
|
||||
async fn run_server() -> Result<(), std::io::Error> {
|
||||
|
||||
// Wow, poem routing gets a bit messy
|
||||
let app = Route::new()
|
||||
.nest("/", StaticFilesEndpoint::new(FILESYSTEM_ROOT)
|
||||
.index_file("index.html")
|
||||
.fallback_to_index()
|
||||
)
|
||||
.at("/download", post(download_url))
|
||||
.at("/download/:path", get(get_file))
|
||||
.with(AddData::new(Arc::new(AppData::default())))
|
||||
.with(poem::middleware::Tracing)
|
||||
;
|
||||
|
||||
Server::new(TcpListener::bind("0.0.0.0:3000"))
|
||||
.run(app)
|
||||
.await
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct Input {
|
||||
url: String,
|
||||
output: OutputFormat,
|
||||
}
|
||||
|
||||
#[derive(Deserialize, Debug)]
|
||||
enum OutputFormat {
|
||||
Video,
|
||||
Audio,
|
||||
}
|
||||
|
||||
impl OutputFormat {
|
||||
fn to_format(&self) -> &str {
|
||||
match self {
|
||||
OutputFormat::Video => "mp4",
|
||||
OutputFormat::Audio => "mp3",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct YtDlpOutput {
|
||||
filename: String,
|
||||
key: String,
|
||||
}
|
||||
|
||||
impl YtDlpOutput {
|
||||
fn from(value: String) -> Self {
|
||||
let x = value.lines();
|
||||
let mut v = Vec::new();
|
||||
x.collect_into(&mut v);
|
||||
Self { filename: v[0].to_owned(), key: v[1].to_owned() }
|
||||
}
|
||||
}
|
||||
|
||||
#[handler]
|
||||
#[instrument(level="trace", skip_all)]
|
||||
async fn download_url(Form(data): Form<Input>, state: Data<&Arc<AppData>>) -> poem::Response {
|
||||
trace!("Download attempt...");
|
||||
// validate url
|
||||
match Url::parse(&data.url) {
|
||||
Ok(url) => {
|
||||
match url.domain() {
|
||||
Some(domain) => {
|
||||
let mut domain: String = domain.to_string();
|
||||
|
||||
let splits: Vec<&str> = domain.split('.').collect();
|
||||
if splits.len() > 2 {
|
||||
domain = [
|
||||
splits[splits.len()-2],
|
||||
".",
|
||||
splits[splits.len()-1]
|
||||
]
|
||||
.iter()
|
||||
.map(|f| *f)
|
||||
.collect();
|
||||
}
|
||||
|
||||
if domain == "youtube.com" || domain == "yout.be" {
|
||||
debug!("Downloading {} to {} format.", data.url, data.output.to_format());
|
||||
|
||||
match Command::new("yt-dlp")
|
||||
.arg("--quiet") // shut up
|
||||
.arg("--no-warnings") // shut up again
|
||||
.arg("--windows-filenames") // windows compatible file name
|
||||
.arg("--no-playlist") // if video or playlist, choose video
|
||||
.args(["--print","after_move:filename,id"]) // output filename
|
||||
.args(["-f", OutputFormat::Video.to_format()]) // format selector
|
||||
.arg(data.url.clone())
|
||||
.output()
|
||||
{
|
||||
Ok(o) => {
|
||||
if let Some(code) = o.status.code() {
|
||||
|
||||
let stderr = o.stderr.iter().map(|f| *f as char).collect::<String>();
|
||||
let stdout = o.stdout.iter().map(|f| *f as char).collect::<String>();
|
||||
|
||||
match code {
|
||||
0 => {
|
||||
info!("Download done.");
|
||||
let mut out = YtDlpOutput::from(stdout);
|
||||
|
||||
// You can potentially do post-processing directly in yt-dlp
|
||||
match data.output {
|
||||
OutputFormat::Video => out.key = format!("video_{}", out.key),
|
||||
OutputFormat::Audio => {
|
||||
trace!("Converting to audio file.");
|
||||
let audio_file = format!("{}.mp3", out.filename);
|
||||
match Command::new("ffmpeg")
|
||||
.args(["-i", &out.filename]) // ffmpeg input
|
||||
.arg(audio_file.clone()) // ffmpeg output
|
||||
.output()
|
||||
{
|
||||
Ok(_) => {
|
||||
out = YtDlpOutput {
|
||||
filename: audio_file,
|
||||
key: format!("audio_{}", out.key),
|
||||
};
|
||||
},
|
||||
Err(_) => error!("Couldn't create shell for ffmpeg."),
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
trace!("Setting state...");
|
||||
state.mapping.write().await.insert(out.key.clone(), out.filename.into());
|
||||
|
||||
let response = poem::Response::builder()
|
||||
.content_type("text/html")
|
||||
.status(StatusCode::UNAUTHORIZED)
|
||||
.body(Body::from_string(format!("<script>window.location = '/download/{}'</script>", out.key)))
|
||||
;
|
||||
|
||||
return response;
|
||||
},
|
||||
2 => {
|
||||
error!("{} STDERR: {}", o.status, stderr);
|
||||
|
||||
}
|
||||
_ => {
|
||||
error!("{}", o.status);
|
||||
error!("{}", stderr);
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
Err(_e) => {
|
||||
error!("Failed to create shell for yt-dlp.");
|
||||
},
|
||||
}
|
||||
}
|
||||
},
|
||||
None => error!("Url didn't have a domain."),
|
||||
}
|
||||
},
|
||||
Err(_) => error!("Couldn't parse url from string."),
|
||||
}
|
||||
poem::Response::builder()
|
||||
.status(StatusCode::INTERNAL_SERVER_ERROR)
|
||||
.body("Internal error")
|
||||
|
||||
}
|
||||
|
||||
|
||||
#[handler]
|
||||
#[instrument(level="trace", skip(state, r))]
|
||||
async fn get_file(Path(request_key): Path<String>,r: &Request, state: Data<&Arc<AppData>>) -> Result<poem::Response, poem::Error> {
|
||||
trace!("Getting: {}", request_key);
|
||||
|
||||
match state.mapping.read().await.get(&request_key) {
|
||||
Some(real_path) => {
|
||||
trace!("Found real path: {:?}", real_path);
|
||||
Ok(StaticFileRequest::from_request_without_body(r)
|
||||
.await?
|
||||
.create_response(real_path, true)?
|
||||
.into_response())
|
||||
},
|
||||
None => {
|
||||
error!("Tried to get {} from the state, but failed.", request_key);
|
||||
Err(poem::Error::from_status(StatusCode::NOT_FOUND))
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user