mirror of
https://github.com/socatlolmeow/drift.git
synced 2026-11-04 11:47:18 +00:00
Add source code
This commit is contained in:
parent
c85d4b5c19
commit
6d386dfe9a
12 changed files with 4900 additions and 0 deletions
320
src/ytmusic.rs
Normal file
320
src/ytmusic.rs
Normal file
|
|
@ -0,0 +1,320 @@
|
|||
use std::process::{Command, Stdio};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use anyhow::{anyhow, bail, Context, Result};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use crate::model::{is_valid_id, parse_clock, Track};
|
||||
|
||||
const ENDPOINT: &str = "https://music.youtube.com/youtubei/v1/search?prettyPrint=false";
|
||||
const SONGS_FILTER: &str = "EgWKAQIIAWoMEA4QChADEAQQCRAF";
|
||||
const UA: &str = "Mozilla/5.0 (X11; Linux x86_64; rv:128.0) Gecko/20100101 Firefox/128.0";
|
||||
const FALLBACK_COUNT: usize = 25;
|
||||
|
||||
pub fn search(query: &str) -> Result<Vec<Track>, String> {
|
||||
match search_innertube(query) {
|
||||
Ok(tracks) if !tracks.is_empty() => Ok(tracks),
|
||||
Ok(_) => ytdlp_search(query).map_err(|e| format!("{e:#}")),
|
||||
Err(primary) => ytdlp_search(query).map_err(|fallback| {
|
||||
format!("YouTube Music search failed ({primary:#}); yt-dlp fallback failed ({fallback:#})")
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn import(url: &str) -> Result<Vec<Track>, String> {
|
||||
if !(url.starts_with("http://") || url.starts_with("https://")) {
|
||||
return Err("import expects a full http(s) URL".into());
|
||||
}
|
||||
ytdlp_entries(url, &[]).map_err(|e| format!("{e:#}"))
|
||||
}
|
||||
|
||||
fn client_version() -> String {
|
||||
let secs = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_secs())
|
||||
.unwrap_or(1_750_000_000);
|
||||
let days = (secs / 86_400) as i64 - 1;
|
||||
let (y, m, d) = civil_from_days(days);
|
||||
format!("1.{y:04}{m:02}{d:02}.01.00")
|
||||
}
|
||||
|
||||
fn civil_from_days(z: i64) -> (i64, i64, i64) {
|
||||
let z = z + 719_468;
|
||||
let era = z.div_euclid(146_097);
|
||||
let doe = z.rem_euclid(146_097);
|
||||
let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365;
|
||||
let y = yoe + era * 400;
|
||||
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
|
||||
let mp = (5 * doy + 2) / 153;
|
||||
let d = doy - (153 * mp + 2) / 5 + 1;
|
||||
let m = if mp < 10 { mp + 3 } else { mp - 9 };
|
||||
(if m <= 2 { y + 1 } else { y }, m, d)
|
||||
}
|
||||
|
||||
fn search_innertube(query: &str) -> Result<Vec<Track>> {
|
||||
let agent = ureq::AgentBuilder::new()
|
||||
.timeout(Duration::from_secs(12))
|
||||
.user_agent(UA)
|
||||
.build();
|
||||
let body = json!({
|
||||
"context": { "client": {
|
||||
"clientName": "WEB_REMIX",
|
||||
"clientVersion": client_version(),
|
||||
"hl": "en",
|
||||
"gl": "US",
|
||||
}},
|
||||
"query": query,
|
||||
"params": SONGS_FILTER,
|
||||
});
|
||||
let resp = agent
|
||||
.post(ENDPOINT)
|
||||
.set("Origin", "https://music.youtube.com")
|
||||
.set("Referer", "https://music.youtube.com/")
|
||||
.set("Cookie", "SOCS=CAI")
|
||||
.send_json(body)
|
||||
.map_err(|e| anyhow!("request failed: {e}"))?;
|
||||
let v: Value = resp.into_json().context("invalid JSON from YouTube Music")?;
|
||||
Ok(parse_search_response(&v))
|
||||
}
|
||||
|
||||
fn collect_items<'a>(v: &'a Value, out: &mut Vec<&'a Value>) {
|
||||
match v {
|
||||
Value::Object(map) => {
|
||||
if let Some(item) = map.get("musicResponsiveListItemRenderer") {
|
||||
out.push(item);
|
||||
} else {
|
||||
for child in map.values() {
|
||||
collect_items(child, out);
|
||||
}
|
||||
}
|
||||
}
|
||||
Value::Array(arr) => arr.iter().for_each(|c| collect_items(c, out)),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
fn find_str<'a>(v: &'a Value, key: &str) -> Option<&'a str> {
|
||||
match v {
|
||||
Value::Object(map) => {
|
||||
if let Some(Value::String(s)) = map.get(key) {
|
||||
return Some(s);
|
||||
}
|
||||
map.values().find_map(|c| find_str(c, key))
|
||||
}
|
||||
Value::Array(arr) => arr.iter().find_map(|c| find_str(c, key)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn column_segments(item: &Value, col: usize) -> Vec<String> {
|
||||
let runs = item
|
||||
.get("flexColumns")
|
||||
.and_then(|c| c.get(col))
|
||||
.and_then(|c| c.pointer("/musicResponsiveListItemFlexColumnRenderer/text/runs"))
|
||||
.and_then(Value::as_array);
|
||||
let mut segments = Vec::new();
|
||||
let mut cur = String::new();
|
||||
for run in runs.into_iter().flatten() {
|
||||
let Some(text) = run.get("text").and_then(Value::as_str) else { continue };
|
||||
if text.trim() == "•" {
|
||||
segments.push(std::mem::take(&mut cur));
|
||||
} else {
|
||||
cur.push_str(text);
|
||||
}
|
||||
}
|
||||
segments.push(cur);
|
||||
segments
|
||||
.into_iter()
|
||||
.map(|s| s.trim().to_string())
|
||||
.filter(|s| !s.is_empty())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn looks_like_duration(s: &str) -> bool {
|
||||
s.contains(':') && s.chars().all(|c| c.is_ascii_digit() || c == ':')
|
||||
}
|
||||
|
||||
fn parse_item(item: &Value) -> Option<Track> {
|
||||
let id = find_str(item, "videoId")?;
|
||||
if !is_valid_id(id) {
|
||||
return None;
|
||||
}
|
||||
let title = column_segments(item, 0).into_iter().next()?;
|
||||
|
||||
let mut meta = column_segments(item, 1);
|
||||
if matches!(meta.first().map(String::as_str), Some("Song" | "Video" | "Episode")) {
|
||||
meta.remove(0);
|
||||
}
|
||||
let duration = match meta.last() {
|
||||
Some(last) if looks_like_duration(last) => {
|
||||
let d = parse_clock(last).map(|s| s as u32);
|
||||
meta.pop();
|
||||
d
|
||||
}
|
||||
_ => None,
|
||||
};
|
||||
let artist = meta.first().cloned().unwrap_or_else(|| "Unknown artist".into());
|
||||
let album = meta
|
||||
.get(1)
|
||||
.filter(|a| !a.contains(" views") && !a.contains(" plays"))
|
||||
.cloned();
|
||||
|
||||
Some(Track {
|
||||
id: id.to_string(),
|
||||
title,
|
||||
artist,
|
||||
album,
|
||||
duration,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn parse_search_response(v: &Value) -> Vec<Track> {
|
||||
let mut items = Vec::new();
|
||||
collect_items(v, &mut items);
|
||||
let mut seen = std::collections::HashSet::new();
|
||||
items
|
||||
.into_iter()
|
||||
.filter_map(parse_item)
|
||||
.filter(|t| seen.insert(t.id.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn ytdlp_search(query: &str) -> Result<Vec<Track>> {
|
||||
ytdlp_entries(&format!("ytsearch{FALLBACK_COUNT}:{query}"), &[])
|
||||
}
|
||||
|
||||
fn ytdlp_entries(target: &str, extra: &[&str]) -> Result<Vec<Track>> {
|
||||
let out = Command::new("yt-dlp")
|
||||
.args([
|
||||
"--flat-playlist",
|
||||
"--dump-json",
|
||||
"--no-warnings",
|
||||
"--ignore-errors",
|
||||
"--socket-timeout",
|
||||
"15",
|
||||
])
|
||||
.args(extra)
|
||||
.arg(target)
|
||||
.stdin(Stdio::null())
|
||||
.output()
|
||||
.map_err(|e| anyhow!("could not run yt-dlp ({e}); is it installed and on PATH?"))?;
|
||||
|
||||
let tracks = parse_ytdlp_lines(&String::from_utf8_lossy(&out.stdout));
|
||||
if tracks.is_empty() && !out.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&out.stderr);
|
||||
let last = stderr
|
||||
.lines()
|
||||
.rev()
|
||||
.find(|l| !l.trim().is_empty())
|
||||
.unwrap_or("unknown error");
|
||||
bail!("yt-dlp failed: {}", last.trim());
|
||||
}
|
||||
Ok(tracks)
|
||||
}
|
||||
|
||||
pub fn parse_ytdlp_lines(text: &str) -> Vec<Track> {
|
||||
let mut seen = std::collections::HashSet::new();
|
||||
text.lines()
|
||||
.filter_map(|line| serde_json::from_str::<Value>(line).ok())
|
||||
.filter_map(|v| {
|
||||
let id = v.get("id")?.as_str()?;
|
||||
if !is_valid_id(id) {
|
||||
return None;
|
||||
}
|
||||
let title = v.get("title")?.as_str()?.trim().to_string();
|
||||
let artist = ["artist", "channel", "uploader"]
|
||||
.iter()
|
||||
.find_map(|k| v.get(*k).and_then(Value::as_str))
|
||||
.unwrap_or("Unknown artist")
|
||||
.trim_end_matches(" - Topic")
|
||||
.to_string();
|
||||
let duration = v.get("duration").and_then(Value::as_f64).map(|d| d as u32);
|
||||
Some(Track {
|
||||
id: id.to_string(),
|
||||
title,
|
||||
artist,
|
||||
album: None,
|
||||
duration,
|
||||
})
|
||||
})
|
||||
.filter(|t| seen.insert(t.id.clone()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn run(text: &str) -> Value {
|
||||
json!({"text": {"runs": [{"text": text}]}})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_song_item() {
|
||||
let resp = json!({"contents": {"shelf": [{"musicResponsiveListItemRenderer": {
|
||||
"playlistItemData": {"videoId": "dQw4w9WgXcQ"},
|
||||
"flexColumns": [
|
||||
{"musicResponsiveListItemFlexColumnRenderer": run("Never Gonna Give You Up")},
|
||||
{"musicResponsiveListItemFlexColumnRenderer": {"text": {"runs": [
|
||||
{"text": "Rick Astley"}, {"text": " • "},
|
||||
{"text": "Whenever You Need Somebody"}, {"text": " • "},
|
||||
{"text": "3:33"}
|
||||
]}}}
|
||||
]
|
||||
}}]}});
|
||||
let tracks = parse_search_response(&resp);
|
||||
assert_eq!(tracks.len(), 1);
|
||||
let t = &tracks[0];
|
||||
assert_eq!(t.id, "dQw4w9WgXcQ");
|
||||
assert_eq!(t.title, "Never Gonna Give You Up");
|
||||
assert_eq!(t.artist, "Rick Astley");
|
||||
assert_eq!(t.album.as_deref(), Some("Whenever You Need Somebody"));
|
||||
assert_eq!(t.duration, Some(213));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_multi_artist_and_type_prefix() {
|
||||
let resp = json!([{"musicResponsiveListItemRenderer": {
|
||||
"overlay": {"watchEndpoint": {"videoId": "abcdefghijk"}},
|
||||
"flexColumns": [
|
||||
{"musicResponsiveListItemFlexColumnRenderer": run("Song")},
|
||||
{"musicResponsiveListItemFlexColumnRenderer": {"text": {"runs": [
|
||||
{"text": "Video"}, {"text": " • "},
|
||||
{"text": "A"}, {"text": " & "}, {"text": "B"}, {"text": " • "},
|
||||
{"text": "1.2M views"}, {"text": " • "}, {"text": "4:01"}
|
||||
]}}}
|
||||
]
|
||||
}}]);
|
||||
let t = &parse_search_response(&resp)[0];
|
||||
assert_eq!(t.artist, "A & B");
|
||||
assert_eq!(t.album, None);
|
||||
assert_eq!(t.duration, Some(241));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skips_items_without_video_id() {
|
||||
let resp = json!([{"musicResponsiveListItemRenderer": {
|
||||
"navigationEndpoint": {"browseId": "UC123"},
|
||||
"flexColumns": [{"musicResponsiveListItemFlexColumnRenderer": run("Artist")}]
|
||||
}}]);
|
||||
assert!(parse_search_response(&resp).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_ytdlp_output() {
|
||||
let text = r#"{"id":"dQw4w9WgXcQ","title":"Song","channel":"Rick - Topic","duration":213.0}
|
||||
not json
|
||||
{"id":"UCxxxxxxxxxxxxxxxxxxxxxx","title":"A channel"}
|
||||
{"id":"dQw4w9WgXcQ","title":"dup"}"#;
|
||||
let t = parse_ytdlp_lines(text);
|
||||
assert_eq!(t.len(), 1);
|
||||
assert_eq!(t[0].artist, "Rick");
|
||||
assert_eq!(t[0].duration, Some(213));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn date_math() {
|
||||
assert_eq!(civil_from_days(0), (1970, 1, 1));
|
||||
assert_eq!(civil_from_days(19_723), (2024, 1, 1));
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue