Commit 46e0854e authored by nitsuga5124's avatar nitsuga5124 🚀 Committed by Taha Hawa

Allow user mentions on impersonate.

	modified:   Cargo.toml
	modified:   src/commands/impersonate.rs
parent 24f8d672
......@@ -15,6 +15,7 @@ bigdecimal = "0.1"
chrono = "0.4"
log = "0.4"
pretty_env_logger = "0.4"
regex = "1.3.6"
[dependencies.serenity]
git = "https://github.com/serenity-rs/serenity.git"
......@@ -29,4 +30,4 @@ features = ["cache", "framework", "builder", "client", "utils", "model", "standa
# opt-level = 'z' # Optimize for size.
opt-level = 3 # Optimize for speed.
lto = true
codegen-units = 1
\ No newline at end of file
codegen-units = 1
......@@ -4,4 +4,11 @@ Master : [![build status](https://gitlab.com/tahahawa/discord-markov-bot/badges/
Discord bot that analyses all users' messages, markov chains them to "impersonate" people
Link to add the "official" instance to your server: https://discordapp.com/oauth2/authorize?client_id=276839250186469376&scope=bot&permissions=67185664
\ No newline at end of file
Link to add the "official" instance to your server: https://discordapp.com/oauth2/authorize?client_id=276839250186469376&scope=bot&permissions=67185664
### How to run:
- Install diesel_cli: `cargo install diesel_cli --no-default-features --features "postgres"`
- Rename `config.yaml.example` to `config.yaml` and fill in the required information.
- Configure the database: `diesel setup --database-url='postgres://localhost/my_db'` (the url is the same as in the configuration file)
- `cargo run --release`
token: YOUR_DISCORD_TOKEN
db: postgres://username:[email protected]:5432/db_name
-- This file should undo anything in `up.sql`
\ No newline at end of file
-- Your SQL goes here
CREATE TABLE guilds(
id bigint NOT NULL PRIMARY KEY,
allow_mention boolean DEFAULT true NOT NULL,
post_limit bigint
);
......@@ -20,13 +20,35 @@ pub fn hivemind(_context: &mut Context, message: &Message, mut args: Args) -> Co
let count: usize = args.single_quoted().unwrap_or(1);
let conn;
{
let conn = {
let mut data = _context.data.write();
let sql_pool = data.get_mut::<Sqlpool>().unwrap().clone();
conn = sql_pool.get().unwrap();
sql_pool.get().unwrap()
};
let guild_results = {
use crate::schema::guilds::dsl::*;
guilds.select(post_limit)
.filter(id.eq(message.guild_id.unwrap().0 as i64))
.limit(1)
.load::<Option<i64>>(&conn)
.expect("Error loading guild information")
};
let limit = {
if guild_results.is_empty() {
None
} else {
guild_results[0]
}
};
if let Some(l) = limit {
if (l as usize) < count {
&message.channel_id.say(&_context, format!("You went past the limit of {} messages", l));
return Ok(());
}
}
let mut chain: Chain<String> = Chain::new();
......
use diesel::dsl::*;
use diesel::prelude::*;
use markov::Chain;
use regex::Regex;
use serenity::{
framework::standard::{macros::command, Args, CommandResult},
model::{channel::Message, id::UserId},
......@@ -21,25 +22,61 @@ enum IdOrUsername {
pub fn impersonate(_context: &mut Context, message: &Message, mut args: Args) -> CommandResult {
debug!("args: {:?}", args);
let conn = {
let mut data = _context.data.write();
let sql_pool = data.get_mut::<Sqlpool>().unwrap().clone();
sql_pool.get().unwrap()
};
let guild_results = {
use crate::schema::guilds::dsl::*;
guilds.select((allow_mention, post_limit))
.filter(id.eq(message.guild_id.unwrap().0 as i64))
.limit(1)
.load::<(bool, Option<i64>)>(&conn)
.expect("Error loading guild information")
};
let (allow_mention, limit) = {
if guild_results.is_empty() {
(true, None)
} else {
guild_results[0]
}
};
let fetch_from = match args.single::<u64>() {
Ok(id) => IdOrUsername::Id(id),
Err(_) => IdOrUsername::Username(args.single_quoted().unwrap_or_else(|_| "".to_owned())),
Err(_) => {
let user = args.single_quoted().unwrap_or_else(|_| "".to_owned());
if user.starts_with("<@") && user.ends_with(">") && allow_mention {
let re = Regex::new("[<@!>]").unwrap();
let user_id = re.replace_all(&user, "").into_owned();
match user_id.parse::<u64>() {
Ok(id) => IdOrUsername::Id(id),
Err(_) => IdOrUsername::Username(user),
}
} else {
IdOrUsername::Username(user)
}
},
};
let count: usize = args.single_quoted().unwrap_or(1);
if let Some(l) = limit {
if (l as usize) < count {
&message.channel_id.say(&_context, format!("You went past the limit of {} messages", l));
return Ok(());
}
}
// let chan = message.channel_id.get().unwrap();
let _ = message.channel_id.broadcast_typing(&_context.http);
let conn;
{
let mut data = _context.data.write();
let sql_pool = data.get_mut::<Sqlpool>().unwrap().clone();
conn = sql_pool.get().unwrap();
}
if let Some(user) = match fetch_from {
IdOrUsername::Id(id) => Some(UserId(id)),
......
......@@ -49,14 +49,80 @@ impl Key for Sqlpool {
type Value = Pool;
}
use schema::guilds;
#[derive(Queryable, Insertable, Debug)]
#[table_name = "guilds"]
pub struct GuildConfig {
pub id: i64,
pub allow_mention: bool,
pub post_limit: Option<i64>,
}
#[command]
fn ping(_ctx: &mut Context, msg: &Message, _args: Args) -> CommandResult {
if let Err(why) = msg.channel_id.say(&_ctx.http, "Pong!") {
fn ping(ctx: &mut Context, msg: &Message, _args: Args) -> CommandResult {
if let Err(why) = msg.channel_id.say(&ctx.http, "Pong!") {
warn!("Error sending message: {:?}", why);
}
Ok(())
}
/// Usage: `~configure limit mentions?`
/// Example: `~configure 100 false`
/// This will allow a maximum of 100 messages for hivemind and impersonate
/// and also will disallow mentions to be used as usernames for impersonate.
///
/// Only the first parameter is required, and setting it to 0 will remove the limit.
/// The second parameter is not required, but it will default to True if not mentioned.
#[command]
#[only_in("guilds")]
#[required_permissions(MANAGE_GUILD)]
#[min_args(1)]
#[aliases(config, conf)]
fn configure(ctx: &mut Context, msg: &Message, mut args: Args) -> CommandResult {
use crate::schema::guilds;
let conn = {
let mut data = ctx.data.write();
let sql_pool = data.get_mut::<Sqlpool>().unwrap().clone();
sql_pool.get().unwrap()
};
// the trait `std::str::FromStr` is not implemented for `std::option::Option<_>`
// :sleep:
let limit = args.single_quoted::<i64>().unwrap_or(0).abs();
let mention = args.single_quoted().unwrap_or(true);
let post_limit = if limit == 0 {
None
} else {
Some(limit)
};
let guild_id = msg.guild_id.unwrap().0 as i64;
let values = GuildConfig {
id: guild_id,
allow_mention: mention,
post_limit: post_limit,
};
diesel::insert_into(guilds::table)
.values(&values)
.on_conflict_do_nothing()
.execute(&conn)?;
diesel::update(guilds::table)
.filter(guilds::id.eq(guild_id))
.set((guilds::allow_mention.eq(mention), guilds::post_limit.eq(post_limit)))
.execute(&conn)?;
&msg.channel_id.say(&ctx, "Success!")?;
Ok(())
}
#[command]
fn stats(_ctx: &mut Context, msg: &Message, _args: Args) -> CommandResult {
let conn;
......@@ -190,7 +256,7 @@ impl EventHandler for Handler {
}
#[group]
#[commands(ping, stats, impersonate, hivemind)]
#[commands(ping, stats, impersonate, hivemind, configure)]
struct General;
......@@ -208,7 +274,7 @@ fn main() {
let manager = diesel::r2d2::ConnectionManager::<PgConnection>::new(dbname.to_string());
let pool = diesel::r2d2::Pool::builder()
.max_size(120)
.max_size(20)
.build(manager)
.unwrap_or_else(|_| panic!("Error connecting to {}", dbname.to_string()));
......
table! {
guilds (id) {
id -> Int8,
allow_mention -> Bool,
post_limit -> Nullable<Int8>,
}
}
table! {
messages (id) {
id -> Int8,
......@@ -9,3 +17,8 @@ table! {
attachment_urls -> Nullable<Array<Text>>,
}
}
allow_tables_to_appear_in_same_query!(
guilds,
messages,
);
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment