Added per-server prefix support.

Signed-off-by: Kamran Mackey <[email protected]>

Author: Kamran Mackey

.gitignore

@@ -1,2 +1,3 @@
 /target
-.env
+.env
+*.sqlite3

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "ellie"
-version = "0.1.9"
+version = "0.2.0"
 authors = ["Kamran Mackey <[email protected]>"]
 publish = false
 build = "build.rs"
@@ -22,6 +22,7 @@ lastfm-rs = { git = "https://github.com/KamranMackey/lastfm-rs" }
 log = "0.4"
 pretty_env_logger = "0.3"
 reqwest = "0.10"
+rusqlite = { git = "https://github.com/KamranMackey/rusqlite", features = ["bundled"] }
 rspotify = { git = "https://github.com/KamranMackey/rspotify" }
 serenity = { git = "https://github.com/serenity-rs/serenity.git", branch = "current" }
 serde = "1.0"

Cargo.toml

@@ -1,3 +1,4 @@
 pub mod help;
 pub mod ping;
+pub mod prefix;
 pub mod version;

src/commands/utils/mod.rs

@@ -0,0 +1,69 @@
+use crate::utilities::database;
+use crate::utilities::database::get_database;
+
+use serenity::framework::standard::macros::command;
+use serenity::framework::standard::Args;
+use serenity::framework::standard::CommandError;
+use serenity::framework::standard::CommandResult;
+use serenity::model::prelude::Message;
+use serenity::prelude::Context;
+
+#[command]
+#[description("Get or set the command prefix for the current server."
+)]
+#[sub_commands(get, set)]
+fn prefix(ctx: &mut Context, message: &Message) -> CommandResult {
+    message
+        .channel_id
+        .send_message(&ctx, move |m| {
+            m.embed(move |e| {
+                e.title("Error: Invalid / No Subcommand Entered!");
+                e.description("Please use subcommand get or set to use this command.");
+                e
+            })
+        })
+        .map_or_else(|e| Err(CommandError(e.to_string())), |_| Ok(()))
+}
+
+#[command]
+#[only_in(guilds)]
+#[owners_only]
+#[description = "Retrieves the command prefix for the current server."]
+pub fn get(ctx: &mut Context, message: &Message) -> CommandResult {
+    let prefix = database::get_prefix(&message.guild_id.unwrap()).unwrap().to_string();
+    let guild = message.guild(&ctx.cache).unwrap();
+    let guild_name = &guild.read().name;
+
+    return message
+        .channel_id
+        .say(&ctx.http, format!("The currently set command prefix for {} is {}.", guild_name, prefix))
+        .map_or_else(|e| Err(CommandError(e.to_string())), |_| Ok(()));
+}
+
+#[command]
+#[only_in(guilds)]
+#[owners_only]
+#[num_args(1)]
+#[description = "Sets the command prefix for the current server."]
+pub fn set(ctx: &mut Context, message: &Message, args: Args) -> CommandResult {
+    let connection = match get_database() {
+        Ok(connection) => connection,
+        Err(_) => return Ok(()),
+    };
+
+    let prefix = args.current().unwrap_or(";");
+    let guild = message.guild(&ctx.cache).unwrap();
+    let guild_id = guild.read().clone().id.as_u64().to_string();
+    let guild_name = guild.read().clone().name;
+
+    let _ = connection.execute(
+        "INSERT OR REPLACE INTO guild_settings (guild_id, guild_name, prefix) values (?1, ?2, ?3)",
+        &[&guild_id, &guild_name, prefix],
+    );
+
+    let _ = message
+        .channel_id
+        .say(&ctx.http, format!("The command prefix for {} has been set to {}.", guild_name, prefix));
+
+    Ok(())
+}

src/commands/utils/prefix.rs

@@ -3,7 +3,7 @@
 //! Ellie is a bot for the Discord chat platform focused on giving users
 //! a powerful set of features, while remaining quick to respond.
 
-pub mod utilities;
+mod utilities;
 mod commands;
 mod listeners;
 
@@ -12,6 +12,7 @@ use commands::info::user::*;
 use commands::music::lastfm::*;
 use commands::music::spotify::commands::spotify::*;
 use commands::utils::help::*;
+use commands::utils::prefix::*;
 use commands::utils::ping::*;
 use commands::utils::version::*;
 
@@ -21,25 +22,28 @@ use log::error;
 
 use listeners::handler::Handler;
 
+use rspotify::spotify::client::Spotify;
+use rspotify::spotify::oauth2::SpotifyClientCredentials;
+
 use serenity::client::Client;
 use serenity::framework::standard::macros::group;
 use serenity::framework::StandardFramework;
 use serenity::framework::standard::DispatchError;
 
-use rspotify::spotify::client::Spotify;
-use rspotify::spotify::oauth2::SpotifyClientCredentials;
-
 use std::collections::HashSet;
 use std::env;
 
+use utilities::database::create_database;
+use utilities::database::get_prefix;
+
 #[group]
 #[description = "Various informational commands."]
 #[commands(user, guild)]
 struct Information;
 
 #[group]
 #[description = "Ellie's selection of utility commands."]
-#[commands(ping, version)]
+#[commands(ping, prefix, version)]
 struct Utilities;
 
 #[group]
@@ -50,8 +54,9 @@ struct Music;
 pub fn main() {
     dotenv().expect("Unable to read / access the .env file!");
 
+    create_database();
+
     let token: String = env::var("DISCORD_TOKEN").expect("Unable to read the bot token.");
-    let prefix: String = env::var("DISCORD_PREFIX").expect("Unable to get the bot prefix.");
 
     let mut client: Client = Client::new(&token, Handler).expect("Error creating client.");
 
@@ -70,10 +75,21 @@ pub fn main() {
         StandardFramework::new()
             .configure(|c| {
                 c.with_whitespace(true)
-                    .prefix(&prefix)
                     .ignore_webhooks(false)
                     .case_insensitivity(true)
                     .owners(owners)
+                    .dynamic_prefix(|_, message| {
+                        let def_prefix: String = env::var("DISCORD_PREFIX").expect("Unable to get the bot prefix.");
+                        if message.is_private() {
+                            return Some(def_prefix.clone().to_string());
+                        }
+                        if let Some(guild_id) = message.guild_id {
+                            let prefix = get_prefix(&guild_id).map_or_else(|_| def_prefix.clone().to_string(), |prefix| prefix);
+                            return Some(prefix);
+                        } else {
+                            return Some(def_prefix.to_string())
+                        }
+                    })
                     .on_mention(Some(bot_id))
             })
             .on_dispatch_error(|ctx, msg, err| match err {

src/main.rs

@@ -0,0 +1,49 @@
+use log::error;
+
+use rusqlite::Connection;
+use rusqlite::NO_PARAMS;
+
+use serenity::model::prelude::GuildId;
+
+use std::error::Error;
+use std::fs::File;
+use std::path::Path;
+
+pub fn create_database() {
+    let database = Path::new("database.sqlite3");
+    if !database.exists() {
+        match File::create(&database) {
+            Ok(_) => (),
+            Err(e) => error!("Failed to create database file: {}", e),
+        }
+    }
+    if let Ok(connection) = Connection::open(&database) {
+        match connection.execute(
+            "CREATE TABLE IF NOT EXISTS guild_settings (
+                guild_id TEXT PRIMARY KEY,
+                guild_name TEXT NOT NULL,
+                prefix TEXT NOT NULL
+            );",
+            NO_PARAMS,
+        ) {
+            Ok(_) => (),
+            Err(e) => {
+                error!("{}", e);
+            }
+        }
+    } else {
+        error!("Could not open connection to database ({})", &database.to_string_lossy());
+    }
+}
+
+pub fn get_database() -> Result<Connection, Box<dyn Error>> {
+    let db = Path::new("database.sqlite3");
+    Ok(Connection::open(db)?)
+}
+
+pub fn get_prefix(guildid: &GuildId) -> Result<String, Box<dyn Error>> {
+    let conn = get_database()?;
+    let mut statement = conn.prepare("SELECT prefix FROM guild_settings WHERE guild_id == ?1;")?;
+    let mut rows = statement.query(&[&guildid.as_u64().to_string()])?;
+    Ok(rows.next()?.ok_or("Guild not found.")?.get(0)?)
+}

src/utilities/database.rs

@@ -1,4 +1,5 @@
 pub mod built_info;
+pub mod database;
 pub mod git_utils;
 
 pub fn format_int(i: isize) -> String {