Skip to main content

bf_bot/
lib.rs

1use serde::Deserialize;
2use serenity::{
3    async_trait,
4    model::{channel::Message, error::Error, gateway::Ready},
5    prelude::*,
6};
7use std::{collections::HashMap, sync::Mutex};
8use tokio::task;
9
10pub struct Handler {
11    user_lock: Mutex<HashMap<u64, HashMap<u64, MapData>>>,
12    config: Config
13}
14
15#[derive(Clone, Deserialize)]
16pub struct Config {
17    pub token: String,
18    pub prefix: String,
19    pub timeout: Option<std::time::Duration>,
20    pub tmppath: Option<std::path::PathBuf>,
21    pub transpile: bool
22}
23
24#[derive(Clone)]
25struct MapData {
26    text: String,
27    botmsg: Message,
28}
29
30impl Handler {
31    pub fn new(config: Config) -> Handler {
32        Handler {
33            user_lock: Mutex::new(HashMap::new()),
34            config
35        }
36    }
37
38    fn add_to_map(&self, chid: u64, uid: u64, content: String, botmsg: Message) {
39        let mut channels = self.user_lock.lock().unwrap();
40        if let Some(user) = channels.get_mut(&chid) {
41            user.insert(
42                uid,
43                MapData {
44                    text: content,
45                    botmsg,
46                },
47            );
48        } else {
49            channels.insert(chid, {
50                let mut map = HashMap::new();
51                map.insert(
52                    uid,
53                    MapData {
54                        text: content,
55                        botmsg,
56                    },
57                );
58                map
59            });
60        }
61    }
62
63    fn get_user_lock(&self, chid: u64, uid: u64) -> Option<MapData> {
64        if let Some(c) = self.user_lock.lock().unwrap().get_mut(&chid) {
65            if let Some(m) = c.remove(&uid) {
66                Some(m)
67            } else {
68                None
69            }
70        } else {
71            None
72        }
73    }
74}
75
76#[async_trait]
77impl EventHandler for Handler {
78    async fn message(&self, ctx: Context, msg: Message) {
79        let mapdata = self.get_user_lock(msg.channel_id.0, msg.author.id.0);
80        let transpile = self.config.transpile;
81        let pl = self.config.prefix.len();
82        let timeout = self.config.timeout;
83        let prog = if let Some(d) = mapdata.clone() {
84            Some(d.text)
85        } else if msg.content.len() > pl + 1 {
86            if msg.content[..pl + 1] == format!("{} ", self.config.prefix) {
87                Some(String::from(&msg.content[2..]))
88            } else {
89                None
90            }
91        } else if msg.content == self.config.prefix && msg.attachments.len() > 0 {
92            match msg.attachments[0].download().await {
93                Ok(chars) => Some(String::from_utf8_lossy(&chars).into_owned()),
94                Err(err) => {
95                    println!("Error downloading attachment: {:?}", err);
96                    None
97                }
98            }
99        } else {
100            None
101        };
102        let input = if let Some(d) = mapdata {
103            if let Err(err) = d.botmsg.delete(&ctx.http).await {
104                println!("Error deleting message: {:?}", err);
105            }
106            Some(msg.content)
107        } else {
108            None
109        };
110        let output = if let Some(prog) = prog {
111            if bf_lib::wants_input(&prog) && input == None {
112                let botmsg = msg
113                    .channel_id
114                    .say(
115                        &ctx.http,
116                        format!(
117                            "Program requires input, next message from {} will be read",
118                            msg.author.name
119                        ),
120                    )
121                    .await
122                    .expect("Error sending message");
123                self.add_to_map(msg.channel_id.0, msg.author.id.0, prog, botmsg);
124                None
125            } else {
126                //let join = task::spawn_blocking(move || bf_lib::Exec::prog(&prog).input(input).timeout(self.config.timeout.clone()).run());
127                let join = task::spawn_blocking(move || 
128                    {
129                        let exec = bf_lib::Exec::prog(&prog).input(input).timeout(timeout);
130                        if transpile {
131                            exec.run()
132                        } else { exec.interpret() }
133                    });
134                let typing = msg.channel_id.start_typing(&ctx.http).unwrap();
135                let o = match join.await.unwrap() {
136                    Ok(ok) => ok,
137                    Err(err) => err.to_string(),
138                };
139                typing.stop();
140                Some(o)
141            }
142        } else {
143            None
144        };
145        if let Some(o) = output {
146            if let Err(e) = msg.channel_id.say(&ctx.http, &o).await {
147                if let serenity::Error::Model(Error::MessageTooLong(_)) = e {
148                    if let Err(e) = msg
149                        .channel_id
150                        .send_files(&ctx.http, vec![(o.as_bytes(), "output.txt")], |m| {
151                            m.content("Program output was too long, sending as file")
152                        })
153                        .await
154                    {
155                        println!("Error sending message: {}", e)
156                    }
157                } else {
158                    println!("Error sending message: {}", e)
159                }
160            };
161        }
162    }
163
164    async fn ready(&self, _: Context, ready: Ready) {
165        println!("{} is connected!", ready.user.name);
166    }
167}