Merge pull request #19 from Mavulp/multibot

Split the bot into a master bot and music bots

2 files changed, 676 insertions, 0 deletions

diff --git a/src/bot/master.rs b/src/bot/master.rs
new file mode 100644
index 0000000..2488064
--- /dev/null
+++ b/src/bot/master.rs
@@ -0,0 +1,268 @@
+use std::collections::HashMap;
+use std::future::Future;
+use std::sync::{Arc, Mutex};
+
+use futures::future::{FutureExt, TryFutureExt};
+use futures01::future::Future as Future01;
+use log::info;
+use rand::{rngs::SmallRng, seq::SliceRandom, SeedableRng};
+use serde::{Deserialize, Serialize};
+use tsclientlib::{ClientId, ConnectOptions, Identity, MessageTarget};
+
+use crate::audio_player::AudioPlayerError;
+use crate::teamspeak::TeamSpeakConnection;
+
+use crate::Args;
+
+use crate::bot::{MusicBot, MusicBotArgs, MusicBotMessage};
+
+pub struct MasterBot {
+    config: Arc<MasterConfig>,
+    music_bots: Arc<Mutex<MusicBots>>,
+    teamspeak: Arc<TeamSpeakConnection>,
+}
+
+struct MusicBots {
+    rng: SmallRng,
+    available_names: Vec<usize>,
+    available_ids: Vec<usize>,
+    connected_bots: HashMap<String, Arc<MusicBot>>,
+}
+
+impl MasterBot {
+    pub async fn new(args: MasterArgs) -> (Arc<Self>, impl Future) {
+        let (tx, mut rx) = tokio02::sync::mpsc::unbounded_channel();
+        let tx = Arc::new(Mutex::new(tx));
+        info!("Starting in TeamSpeak mode");
+
+        let mut con_config = ConnectOptions::new(args.address.clone())
+            .version(tsclientlib::Version::Linux_3_3_2)
+            .name(args.master_name.clone())
+            .identity(args.id)
+            .log_commands(args.verbose >= 1)
+            .log_packets(args.verbose >= 2)
+            .log_udp_packets(args.verbose >= 3);
+
+        if let Some(channel) = args.channel {
+            con_config = con_config.channel(channel);
+        }
+
+        let connection = Arc::new(
+            TeamSpeakConnection::new(tx.clone(), con_config)
+                .await
+                .unwrap(),
+        );
+
+        let config = Arc::new(MasterConfig {
+            master_name: args.master_name,
+            address: args.address,
+            names: args.names,
+            ids: args.ids,
+            local: args.local,
+            verbose: args.verbose,
+        });
+
+        let name_count = config.names.len();
+        let id_count = config.ids.len();
+
+        let music_bots = Arc::new(Mutex::new(MusicBots {
+            rng: SmallRng::from_entropy(),
+            available_names: (0..name_count).collect(),
+            available_ids: (0..id_count).collect(),
+            connected_bots: HashMap::new(),
+        }));
+
+        let bot = Arc::new(Self {
+            config,
+            music_bots,
+            teamspeak: connection,
+        });
+
+        bot.teamspeak
+            .set_description("Poke me if you want a music bot!");
+
+        let cbot = bot.clone();
+        let msg_loop = async move {
+            loop {
+                while let Some(msg) = rx.recv().await {
+                    cbot.on_message(msg).await.unwrap();
+                }
+            }
+        };
+
+        (bot, msg_loop)
+    }
+
+    fn build_bot_args_for(&self, id: ClientId) -> Option<MusicBotArgs> {
+        let channel = self
+            .teamspeak
+            .channel_of_user(id)
+            .expect("Can find poke sender");
+
+        if channel == self.teamspeak.my_channel() {
+            self.teamspeak.send_message_to_user(
+                id,
+                &format!(
+                    "Joining the channel of \"{}\" is not allowed",
+                    self.config.master_name
+                ),
+            );
+            return None;
+        }
+
+        let MusicBots {
+            ref mut rng,
+            ref mut available_names,
+            ref mut available_ids,
+            ref connected_bots,
+        } = &mut *self.music_bots.lock().expect("Mutex was not poisoned");
+
+        for (_, bot) in connected_bots {
+            if bot.my_channel() == channel {
+                self.teamspeak.send_message_to_user(
+                    id,
+                    &format!(
+                        "\"{}\" is already in this channel. \
+                         Multiple bots in one channel are not allowed.",
+                        bot.name()
+                    ),
+                );
+                return None;
+            }
+        }
+
+        let channel_path = self
+            .teamspeak
+            .channel_path_of_user(id)
+            .expect("can find poke sender");
+
+        available_names.shuffle(rng);
+        let name_index = match available_names.pop() {
+            Some(v) => v,
+            None => {
+                self.teamspeak
+                    .send_message_to_user(id, "Out of names. Too many bots are already connected!");
+                return None;
+            }
+        };
+        let name = self.config.names[name_index].clone();
+
+        available_ids.shuffle(rng);
+        let id_index = match available_ids.pop() {
+            Some(v) => v,
+            None => {
+                self.teamspeak.send_message_to_user(
+                    id,
+                    "Out of identities. Too many bots are already connected!",
+                );
+                return None;
+            }
+        };
+
+        let id = self.config.ids[id_index].clone();
+
+        let cmusic_bots = self.music_bots.clone();
+        let disconnect_cb = Box::new(move |n, name_index, id_index| {
+            let mut music_bots = cmusic_bots.lock().expect("Mutex was not poisoned");
+            music_bots.connected_bots.remove(&n);
+            music_bots.available_names.push(name_index);
+            music_bots.available_ids.push(id_index);
+        });
+
+        info!("Connecting to {} on {}", channel_path, self.config.address);
+
+        Some(MusicBotArgs {
+            name,
+            name_index,
+            id_index,
+            local: self.config.local,
+            address: self.config.address.clone(),
+            id,
+            channel: channel_path,
+            verbose: self.config.verbose,
+            disconnect_cb,
+        })
+    }
+
+    async fn spawn_bot_for(&self, id: ClientId) {
+        if let Some(bot_args) = self.build_bot_args_for(id) {
+            let (bot, fut) = MusicBot::new(bot_args).await;
+            tokio::spawn(fut.unit_error().boxed().compat().map(|_| ()));
+            let mut music_bots = self.music_bots.lock().expect("Mutex was not poisoned");
+            music_bots
+                .connected_bots
+                .insert(bot.name().to_string(), bot);
+        }
+    }
+
+    async fn on_message(&self, message: MusicBotMessage) -> Result<(), AudioPlayerError> {
+        if let MusicBotMessage::TextMessage(message) = message {
+            if let MessageTarget::Poke(who) = message.target {
+                info!("Poked by {}, creating bot for their channel", who);
+                self.spawn_bot_for(who).await;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct MasterArgs {
+    #[serde(default = "default_name")]
+    pub master_name: String,
+    #[serde(default = "default_local")]
+    pub local: bool,
+    pub address: String,
+    pub channel: Option<String>,
+    #[serde(default = "default_verbose")]
+    pub verbose: u8,
+    pub names: Vec<String>,
+    pub id: Identity,
+    pub ids: Vec<Identity>,
+}
+
+fn default_name() -> String {
+    String::from("PokeBot")
+}
+
+fn default_local() -> bool {
+    false
+}
+
+fn default_verbose() -> u8 {
+    0
+}
+
+impl MasterArgs {
+    pub fn merge(self, args: Args) -> Self {
+        let address = args.address.unwrap_or(self.address);
+        let local = args.local || self.local;
+        let channel = args.master_channel.or(self.channel);
+        let verbose = if args.verbose > 0 {
+            args.verbose
+        } else {
+            self.verbose
+        };
+
+        Self {
+            master_name: self.master_name,
+            names: self.names,
+            ids: self.ids,
+            local,
+            address,
+            id: self.id,
+            channel,
+            verbose,
+        }
+    }
+}
+
+pub struct MasterConfig {
+    pub master_name: String,
+    pub address: String,
+    pub names: Vec<String>,
+    pub ids: Vec<Identity>,
+    pub local: bool,
+    pub verbose: u8,
+}
diff --git a/src/bot/music.rs b/src/bot/music.rs
new file mode 100644
index 0000000..dee1514
--- /dev/null
+++ b/src/bot/music.rs
@@ -0,0 +1,408 @@
+use std::future::Future;
+use std::io::BufRead;
+use std::sync::{Arc, Mutex};
+use std::thread;
+
+use log::{debug, info};
+use structopt::StructOpt;
+use tokio02::sync::mpsc::UnboundedSender;
+use tsclientlib::{data, ChannelId, ClientId, ConnectOptions, Identity, Invoker, MessageTarget};
+
+use crate::audio_player::{AudioPlayer, AudioPlayerError, PollResult};
+use crate::command::Command;
+use crate::playlist::Playlist;
+use crate::teamspeak::TeamSpeakConnection;
+use crate::youtube_dl::AudioMetadata;
+
+#[derive(Debug)]
+pub struct Message {
+    pub target: MessageTarget,
+    pub invoker: Invoker,
+    pub text: String,
+}
+
+#[derive(Debug, PartialEq, Eq)]
+pub enum State {
+    Playing,
+    Paused,
+    Stopped,
+    EndOfStream,
+}
+
+#[derive(Debug)]
+pub enum MusicBotMessage {
+    TextMessage(Message),
+    ClientChannel {
+        client: ClientId,
+        old_channel: ChannelId,
+    },
+    ClientDisconnected {
+        id: ClientId,
+        client: data::Client,
+    },
+    StateChange(State),
+    Quit(String),
+}
+
+pub struct MusicBot {
+    name: String,
+    player: Arc<AudioPlayer>,
+    teamspeak: Option<Arc<TeamSpeakConnection>>,
+    playlist: Arc<Mutex<Playlist>>,
+    state: Arc<Mutex<State>>,
+}
+
+pub struct MusicBotArgs {
+    pub name: String,
+    pub name_index: usize,
+    pub id_index: usize,
+    pub local: bool,
+    pub address: String,
+    pub id: Identity,
+    pub channel: String,
+    pub verbose: u8,
+    pub disconnect_cb: Box<dyn FnMut(String, usize, usize) + Send + Sync>,
+}
+
+impl MusicBot {
+    pub async fn new(args: MusicBotArgs) -> (Arc<Self>, impl Future) {
+        let (tx, mut rx) = tokio02::sync::mpsc::unbounded_channel();
+        let tx = Arc::new(Mutex::new(tx));
+        let (player, connection) = if args.local {
+            info!("Starting in CLI mode");
+            let audio_player = AudioPlayer::new(tx.clone(), None).unwrap();
+
+            (audio_player, None)
+        } else {
+            info!("Starting in TeamSpeak mode");
+
+            let con_config = ConnectOptions::new(args.address)
+                .version(tsclientlib::Version::Linux_3_3_2)
+                .name(format!("🎵 {}", args.name))
+                .identity(args.id)
+                .log_commands(args.verbose >= 1)
+                .log_packets(args.verbose >= 2)
+                .log_udp_packets(args.verbose >= 3)
+                .channel(args.channel);
+
+            let connection = Arc::new(
+                TeamSpeakConnection::new(tx.clone(), con_config)
+                    .await
+                    .unwrap(),
+            );
+            let cconnection = connection.clone();
+            let audio_player = AudioPlayer::new(
+                tx.clone(),
+                Some(Box::new(move |samples| {
+                    cconnection.send_audio_packet(samples);
+                })),
+            )
+            .unwrap();
+
+            (audio_player, Some(connection))
+        };
+
+        player.set_volume(0.5).unwrap();
+        let player = Arc::new(player);
+        let playlist = Arc::new(Mutex::new(Playlist::new()));
+
+        spawn_gstreamer_thread(player.clone(), tx.clone());
+
+        if args.local {
+            spawn_stdin_reader(tx);
+        }
+
+        let bot = Arc::new(Self {
+            name: args.name.clone(),
+            player,
+            teamspeak: connection,
+            playlist,
+            state: Arc::new(Mutex::new(State::Stopped)),
+        });
+
+        let cbot = bot.clone();
+        let mut disconnect_cb = args.disconnect_cb;
+        let name = args.name;
+        let name_index = args.name_index;
+        let id_index = args.id_index;
+        let msg_loop = async move {
+            'outer: loop {
+                while let Some(msg) = rx.recv().await {
+                    if let MusicBotMessage::Quit(reason) = msg {
+                        cbot.with_teamspeak(|ts| ts.disconnect(&reason));
+                        disconnect_cb(name, name_index, id_index);
+                        break 'outer;
+                    }
+                    cbot.on_message(msg).await.unwrap();
+                }
+            }
+            debug!("Left message loop");
+        };
+
+        (bot, msg_loop)
+    }
+
+    #[inline(always)]
+    fn with_teamspeak<F: Fn(&TeamSpeakConnection)>(&self, func: F) {
+        if let Some(ts) = &self.teamspeak {
+            func(&ts);
+        }
+    }
+
+    fn start_playing_audio(&self, metadata: AudioMetadata) {
+        if let Some(title) = metadata.title {
+            self.send_message(&format!("Playing '{}'", title));
+            self.set_description(&format!("Currently playing '{}'", title));
+        } else {
+            self.send_message("Playing unknown title");
+            self.set_description("Currently playing");
+        }
+        self.player.reset().unwrap();
+        self.player.set_source_url(metadata.url).unwrap();
+        self.player.play().unwrap();
+    }
+
+    pub async fn add_audio(&self, url: String) {
+        match crate::youtube_dl::get_audio_download_url(url).await {
+            Ok(metadata) => {
+                info!("Found audio url: {}", metadata.url);
+
+                let mut playlist = self.playlist.lock().expect("Mutex was not poisoned");
+                playlist.push(metadata.clone());
+
+                if !self.player.is_started() {
+                    if let Some(request) = playlist.pop() {
+                        self.start_playing_audio(request);
+                    }
+                } else {
+                    if let Some(title) = metadata.title {
+                        self.send_message(&format!("Added '{}' to playlist", title));
+                    } else {
+                        self.send_message("Added to playlist");
+                    }
+                }
+            }
+            Err(e) => {
+                info!("Failed to find audio url: {}", e);
+
+                self.send_message(&format!("Failed to find url: {}", e));
+            }
+        }
+    }
+
+    pub fn name(&self) -> &str {
+        &self.name
+    }
+
+    pub fn my_channel(&self) -> ChannelId {
+        self.teamspeak
+            .as_ref()
+            .map(|ts| ts.my_channel())
+            .expect("my_channel needs ts")
+    }
+
+    fn user_count(&self, channel: ChannelId) -> u32 {
+        self.teamspeak
+            .as_ref()
+            .map(|ts| ts.user_count(channel))
+            .expect("user_count needs ts")
+    }
+
+    fn send_message(&self, text: &str) {
+        debug!("Sending message to TeamSpeak: {}", text);
+
+        self.with_teamspeak(|ts| ts.send_message_to_channel(text));
+    }
+
+    fn set_nickname(&self, name: &str) {
+        info!("Setting TeamsSpeak nickname to {}", name);
+
+        self.with_teamspeak(|ts| ts.set_nickname(name));
+    }
+
+    fn set_description(&self, desc: &str) {
+        info!("Setting TeamsSpeak description to {}", desc);
+
+        self.with_teamspeak(|ts| ts.set_description(desc));
+    }
+
+    async fn on_text(&self, message: Message) -> Result<(), AudioPlayerError> {
+        let msg = message.text;
+        if msg.starts_with("!") {
+            let tokens = msg[1..].split_whitespace().collect::<Vec<_>>();
+
+            match Command::from_iter_safe(&tokens) {
+                Ok(args) => self.on_command(args).await?,
+                Err(e) if e.kind == structopt::clap::ErrorKind::HelpDisplayed => {
+                    self.send_message(&format!("\n{}", e.message));
+                }
+                _ => (),
+            }
+        }
+
+        Ok(())
+    }
+
+    async fn on_command(&self, command: Command) -> Result<(), AudioPlayerError> {
+        match command {
+            Command::Play => {
+                let playlist = self.playlist.lock().expect("Mutex was not poisoned");
+
+                if !self.player.is_started() {
+                    if !playlist.is_empty() {
+                        self.player.stop_current()?;
+                    }
+                } else {
+                    self.player.play()?;
+                }
+            }
+            Command::Add { url } => {
+                // strip bbcode tags from url
+                let url = url.replace("[URL]", "").replace("[/URL]", "");
+
+                self.add_audio(url.to_string()).await;
+            }
+            Command::Pause => {
+                self.player.pause()?;
+            }
+            Command::Stop => {
+                self.player.reset()?;
+            }
+            Command::Next => {
+                let playlist = self.playlist.lock().expect("Mutex was not poisoned");
+                if !playlist.is_empty() {
+                    info!("Skipping to next track");
+                    self.player.stop_current()?;
+                } else {
+                    info!("Playlist empty, cannot skip");
+                    self.player.reset()?;
+                }
+            }
+            Command::Clear => {
+                self.playlist
+                    .lock()
+                    .expect("Mutex was not poisoned")
+                    .clear();
+            }
+            Command::Volume { percent: volume } => {
+                let volume = volume.max(0.0).min(100.0) * 0.01;
+                self.player.set_volume(volume)?;
+            }
+            Command::Leave => {
+                self.quit(String::from("Leaving"));
+            }
+        }
+
+        Ok(())
+    }
+
+    fn on_state(&self, state: State) -> Result<(), AudioPlayerError> {
+        let mut current_state = self.state.lock().unwrap();
+        if *current_state != state {
+            match state {
+                State::Playing => {
+                    self.set_nickname(&format!("🎵 {} - Playing", self.name));
+                }
+                State::Paused => {
+                    self.set_nickname(&format!("🎵 {} - Paused", self.name));
+                }
+                State::Stopped => {
+                    self.set_nickname(&format!("🎵 {}", self.name));
+                    self.set_description("");
+                }
+                State::EndOfStream => {
+                    let next_track = self.playlist.lock().expect("Mutex was not poisoned").pop();
+                    if let Some(request) = next_track {
+                        info!("Advancing playlist");
+
+                        self.start_playing_audio(request);
+                    } else {
+                        self.set_nickname(&format!("🎵 {}", self.name));
+                        self.set_description("");
+                    }
+                }
+            }
+        }
+
+        *current_state = state;
+
+        Ok(())
+    }
+
+    async fn on_message(&self, message: MusicBotMessage) -> Result<(), AudioPlayerError> {
+        match message {
+            MusicBotMessage::TextMessage(message) => {
+                if MessageTarget::Channel == message.target {
+                    self.on_text(message).await?;
+                }
+            }
+            MusicBotMessage::ClientChannel {
+                client: _,
+                old_channel,
+            } => {
+                self.on_client_left_channel(old_channel);
+            }
+            MusicBotMessage::ClientDisconnected { id: _, client } => {
+                let old_channel = client.channel;
+                self.on_client_left_channel(old_channel);
+            }
+            MusicBotMessage::StateChange(state) => {
+                self.on_state(state)?;
+            }
+            MusicBotMessage::Quit(_) => (),
+        }
+
+        Ok(())
+    }
+
+    fn on_client_left_channel(&self, old_channel: ChannelId) {
+        let my_channel = self.my_channel();
+        if old_channel == my_channel && self.user_count(my_channel) <= 1 {
+            self.quit(String::from("Channel is empty"));
+        }
+    }
+
+    pub fn quit(&self, reason: String) {
+        self.player.quit(reason);
+    }
+}
+
+fn spawn_stdin_reader(tx: Arc<Mutex<UnboundedSender<MusicBotMessage>>>) {
+    debug!("Spawning stdin reader thread");
+    thread::spawn(move || {
+        let stdin = ::std::io::stdin();
+        let lock = stdin.lock();
+        for line in lock.lines() {
+            let line = line.unwrap();
+
+            let message = MusicBotMessage::TextMessage(Message {
+                target: MessageTarget::Channel,
+                invoker: Invoker {
+                    name: String::from("stdin"),
+                    id: ClientId(0),
+                    uid: None,
+                },
+                text: line,
+            });
+
+            let tx = tx.lock().unwrap();
+            tx.send(message).unwrap();
+        }
+    });
+}
+
+fn spawn_gstreamer_thread(
+    player: Arc<AudioPlayer>,
+    tx: Arc<Mutex<UnboundedSender<MusicBotMessage>>>,
+) {
+    thread::spawn(move || loop {
+        if player.poll() == PollResult::Quit {
+            break;
+        }
+
+        tx.lock()
+            .unwrap()
+            .send(MusicBotMessage::StateChange(State::EndOfStream))
+            .unwrap();
+    });
+}