rapblock/src/main.rs

#[macro_use] extern crate serde_derive;
//#[macro_use] extern crate lazy_static;
#[macro_use] extern crate log;
#[macro_use] extern crate failure;

mod brainz;

use mpris::{PlayerFinder,Event};
use failure::Error;
use std::time::Duration;
use std::collections::HashSet;
use std::env;
use std::path::Path;
use std::fs::File;
use std::io::Read;
use std::collections::HashMap;

#[derive(Serialize, Deserialize, Debug)]
#[serde(default)]
pub struct BlacklistConf {
    /// Literal tags to reject
    pub tag: Vec<String>,
    /// Tag sub-strings to reject (must be a whole word)
    pub tag_partial: Vec<String>,
    /// Artists to reject
    pub artist: Vec<String>,
}

impl Default for BlacklistConf {
    fn default() -> Self {
        Self {
            tag: vec![],
            artist: vec![],
            tag_partial: vec![
                "hip-hop".to_owned(),
                "hip hop".to_owned(),
                "hiphop".to_owned(),
                "rap".to_owned()
            ],
        }
    }
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(default)]
pub struct WhiteList {
    /// Tags to allow despite e.g. a substring match
    pub tag: Vec<String>,
    /// Artists to allow
    pub artist: Vec<String>,
}

impl Default for WhiteList {
    fn default() -> Self {
        Self {
            tag: vec![],
            artist: vec![]
        }
    }
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(default)]
pub struct Config {
    /// Logging - trace, debug, (info), warning, error
    pub logging: String,
    /// Players to handle (empty = all)
    pub allowed_players: Vec<String>,
    /// Blacklists
    pub blacklist: BlacklistConf,
    /// Whitelist (overrides blacklist)
    pub whitelist: WhiteList,
    /// Min MusicBrainz search score for artist look-up
    pub artist_min_score : i32,
    /// Max nr of artists to check per track
    pub max_artists_per_track : u64,
    /// Interval in which the daemon probes DBUS for open MPRIS channels
    pub player_find_interval_ms: u64,
    /// Delay after a skip or allow, e.g. to prevent infinite skip chain when someone starts a rap playlist
    pub cooldown_ms: u64,
    /// Music Brainz API access timeout
    pub api_timeout_ms: u64,
    /// Allow playing songs from artists we couldn't verify
    pub allow_by_default: bool,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            logging: "info".to_owned(),
            allowed_players: vec![],
            blacklist: BlacklistConf::default(),
            whitelist: WhiteList::default(),
            artist_min_score: 95,
            max_artists_per_track: 3,
            player_find_interval_ms: 2500,
            cooldown_ms: 1000,
            api_timeout_ms: 2000,
            allow_by_default: true
        }
    }
}

const CONFIG_FILE: &'static str = "rapblock.json";

fn load_config() -> Result<Config, Error> {
    let config_file_path = env::current_dir()?.join(CONFIG_FILE);
    let buf = read_file(config_file_path)?;

    let config : Config = serde_json::from_str(&buf)?;

    // Validations
    match config.logging.as_ref() {
        "info" | "debug" | "trace" | "warning" | "error" => (),
        _ => bail!("Invalid value for \"logging\""),
    };

    Ok(config)
}

pub fn read_file<P: AsRef<Path>>(path: P) -> Result<String, Error> {
    let path = path.as_ref();
    let mut file = File::open(path)?;

    let mut buf = String::new();
    file.read_to_string(&mut buf)?;
    Ok(buf)
}

fn main() -> Result<(), Error> {
    let config = match load_config() {
        Ok(c) => c,
        Err(e) => {
            eprintln!("Could not load config from \"{}\": {}", CONFIG_FILE, e);
            Config::default()
        }
    };

    let env = env_logger::Env::default().default_filter_or(&config.logging);
    env_logger::Builder::from_env(env).init();

    debug!("Loaded config:\n{}", serde_json::to_string_pretty(&config)?);

    let mut artist_cache = HashMap::<String, bool>::new();

    info!("Waiting for players...");

    'main_loop:
    loop {
        // XXX this picks the first player, which isn't always ideal - see mpris/src/find.rs
        let player = PlayerFinder::new()
            .expect("Could not connect to D-Bus")
            .find_active();

        if let Ok(player) = player {
            let player_name = player.bus_name().to_string();
            if !config.allowed_players.is_empty() && !config.allowed_players.contains(&player_name) {
                debug!("Ignoring player {}", player_name);
                ::std::thread::sleep(Duration::from_millis(config.player_find_interval_ms));
                continue 'main_loop;
            }

            info!("Connected to player: {}{}", player_name, player.unique_name());

            let events = player.events();
            if events.is_err() {
                error!("Could not start event stream!");
                // add a delay so we don't run too hot here
                ::std::thread::sleep(Duration::from_millis(config.player_find_interval_ms));
                continue 'main_loop;
            }

            'event_loop:
            for event in events.unwrap() {
                match event {
                    Ok(event) => {
                        debug!("MPRIS event: {:#?}", event);
                        match &event {
                            Event::PlayerShutDown => {
                                info!("Player shut down");
                                break 'event_loop;
                            },
                            Event::TrackChanged(metadata) => {
                                let title = metadata.title().unwrap_or("");
                                info!("--- new track : {} ---", title);
                                debug!("{:#?}", event);

                                if title.is_empty() {
                                    warn!("!!! Spotify is giving us garbage - empty metadata struct !!!");
                                    // wait for next event
                                    continue 'event_loop;
                                }

                                let mut artists = HashSet::new();

                                if let Some(aa) = metadata.artists() {
                                    for a in aa { artists.insert(a); }
                                }
                                if let Some(aa) = metadata.album_artists() {
                                    for a in aa { artists.insert(a); }
                                }

                                let mut skip = false;
                                let mut confidence = false;
                                'artists_loop:
                                for (an, a) in artists.iter().take(config.max_artists_per_track as usize).enumerate() {
                                    info!("Checking artist #{}: {}", an+1, a);
                                    if let Some(resolution) = artist_cache.get(a.as_str()) {
                                        confidence = true;
                                        if !resolution {
                                            info!("~ result cached: BAD");
                                            skip = true;
                                            break 'artists_loop;
                                        }
                                        info!("~ result cached: GOOD");
                                        continue 'artists_loop;
                                    }

                                    if config.whitelist.artist.contains(a) {
                                        info!("+ Whitelisted artist!");
                                        // there may be other co-artists that spoil the song -> don't break yet
                                        artist_cache.insert(a.to_string(), true);
                                        confidence = true;
                                        continue 'artists_loop;
                                    }

                                    if config.blacklist.artist.contains(a) {
                                        info!("- Blacklisted artist!");
                                        skip = true;
                                        artist_cache.insert(a.to_string(), false);
                                        confidence = true;
                                        break 'artists_loop;
                                    }

                                    let verdict = brainz::check_artist(&config, &a);
                                    match verdict {
                                        Ok(allow) => {
                                            confidence = true;
                                            artist_cache.insert(a.to_string(), allow);
                                            if allow {
                                                info!("Artist passed");
                                            } else {
                                                skip = true;
                                                break 'artists_loop;
                                            }
                                        },
                                        Err(e) => {
                                            warn!("Something went wrong: {}", e);
                                            // probably no tags, or not found - use the default action
                                            artist_cache.insert(a.to_string(), config.allow_by_default);
                                        }
                                    }
                                }

                                if skip || (!confidence && !config.allow_by_default) {
                                    info!(">>>>>> SKIP : {} >>>>>>\n", title);
                                    if player.next().is_err() {
                                        break 'event_loop;
                                    }
                                } else {
                                    info!("Let it play...\n");
                                }

                                ::std::thread::sleep(Duration::from_millis(config.cooldown_ms));
                            },
                            _ => {
                                debug!("Event not handled.");
                            }
                        }
                    },
                    Err(err) => {
                        error!("D-Bus error: {}. Aborting.", err);
                        break 'event_loop;
                    }
                }
            }
            info!("Event stream ended - player likely shut down");
        } else {
            debug!("No player found, waiting...");
        }

        ::std::thread::sleep(Duration::from_millis(config.player_find_interval_ms));
    }
}