ByteDice/ByteDiceAssistant
1
0
Fork 0
Code Issues Pull Requests Actions Releases Wiki Activity

put ASSISTANT_OWNERS in the ctx.data()

Browse Source
This commit is contained in:
ByteDice
2025-03-02 12:18:49 +01:00
parent 2774b56e30
commit ab02ac1cf6
4 changed files with 29 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+2 -2
View File
@@ -38,10 +38,10 @@ You can install Python modules by running `$ pip install {module}` or `$ python
| `ASSISTANT_R_NAME` | The username of the Reddit bot/account. |
| `ASSISTANT_R_PASS` | The password for the Reddit bot/account. |
| `ASSISTANT_OWNERS` | (OPTIONAL) A list of Discord user IDs that "own" the bot. Separate each ID with a single comma and **no** spaces. This will allow the specified user IDs to run root commands such as `/stop`, it will also DM these users when *certain* errors occur. |
| `ASSISTANT_BK_WEEK_MODS` | (OPTIONAL) Same format as `ASSISTANT_OWNERS` but for |
| `ASSISTANT_BK_WEEK_MODS` | (OPTIONAL) Same format as `ASSISTANT_OWNERS` but for people who are allowed to use the `/bk_week` commands. |
### Required permissions:
**These are automatically set if you use the [official invite link](https://discord.com/oauth2/authorize?client_id=1212127255795335208&permissions=84992&integration_type=0&scope=bot).**
**These are automatically set if you use the [official invite link](https://discord.com/oauth2/authorize?client_id=1212127255795335208&permissions=84992&integration_type=0&scope=bot) or an invite link with the permissions integer set to `84992`.**
* Send Messages
* Read Message History
* View Channels
src/main.rs
+16 -14
View File
@@ -58,6 +58,7 @@ type Context<'a> = poise::Context<'a, Data, Error>;
struct Data {
owners: Vec<u64>,
ball_prompts: [Vec<String>; 2],
reddit_data: Mutex<Option<Value>>,
discord_data: Mutex<Option<Value>>,
@@ -74,6 +75,13 @@ async fn main() {
let args = <Args as clap::Parser>::parse();
let args_str = serde_json::to_string(&args).expect("Error serializing args to JSON");
let own_env = std::env::var("ASSISTANT_OWNERS").expect("Missing ASSISTANT_OWNERS env var!");
let own_vec_str: Vec<String> = own_env.split(",").map(String::from).collect();
let own_vec_u64: Vec<u64> = own_vec_str
.iter()
.filter_map(|s| Some(s.parse::<u64>().expect("Failed to parse ASSISTANT_OWNERS. Invalid syntax.")))
.collect();
if args.test { println!("----- USING TEST BOT -----"); }
if args.dev { println!("----- DEV MODE ENABLED -----"); }
if args.dev && args.wipe { println!("----- \"DON'T WORRY ABOUT IT\" MODE ENABLED -----"); }
@@ -88,7 +96,7 @@ async fn main() {
else if args.rs && ! args.py {
println!("----- RUST ONLY MODE -----");
rs_println!("ARGS: {}", args_str);
start(args).await;
start(args, own_vec_u64.clone()).await;
process::exit(0);
}
else if args.py && args.rs {
@@ -103,8 +111,8 @@ async fn main() {
let rust = thread::spawn(move || {
rt.block_on(async {
websocket::start(rust_args.clone()).await;
start(rust_args).await;
websocket::start(rust_args.clone(), own_vec_u64.clone()).await;
start(rust_args, own_vec_u64).await;
});
});
@@ -125,8 +133,8 @@ async fn main() {
}
async fn start(args: Args) {
let data = gen_data(args.clone()).await;
async fn start(args: Args, owners: Vec<u64>) {
let data = gen_data(args.clone(), owners).await;
let mut bot = gen_bot(data, args).await;
rs_println!("Starting bot...");
@@ -134,7 +142,7 @@ async fn start(args: Args) {
}
async fn gen_data(args: Args) -> Data {
async fn gen_data(args: Args, owners: Vec<u64>) -> Data {
let ball_classic_str = std::fs::read_to_string("./data/8-ball_classic.txt").unwrap();
let ball_quirk_str = std::fs::read_to_string("./data/8-ball_quirky.txt").unwrap();
let bk_mods_str = std::fs::read_to_string("./data/bk_mods.json").unwrap();
@@ -144,6 +152,7 @@ async fn gen_data(args: Args) -> Data {
let bk_mods: Value = serde_json::from_str(&bk_mods_str).unwrap();
let data = Data {
owners,
ball_prompts: [ball_classic, ball_quirk],
bk_mods_json: bk_mods,
reddit_data: None.into(),
@@ -174,14 +183,7 @@ async fn gen_bot(data: Data, args: Args) -> Client {
let token_end_len = token[peek_len..].len();
rs_println!("Token: {}{}", token_peek, "*".repeat(token_end_len));
let own_env = std::env::var("ASSISTANT_OWNERS").expect("Missing ASSISTANT_OWNERS env var!");
let own_vec_str: Vec<String> = own_env.split(",").map(String::from).collect();
let own_vec_u64: Vec<u64> = own_vec_str
.iter()
.filter_map(|s| Some(s.parse::<u64>().expect("Failed to parse ASSISTANT_OWNERS. Invalid syntax.")))
.collect();
let own: HashSet<UserId> = own_vec_u64.into_iter().map(UserId::from).collect();
let own: HashSet<UserId> = data.owners.clone().into_iter().map(UserId::from).collect();
let framework = poise::Framework::builder()
.options(poise::FrameworkOptions {
src/messages.rs
+5 -13
View File
@@ -178,25 +178,17 @@ pub async fn http_edit_msg(
}
pub async fn send_dm(msg: String, args: Args) {
let uids = env::var("ASSISTANT_OWNERS");
if uids.is_err() { return; }
let uids_split: Vec<String> = uids.unwrap()
.split(',')
.map(|s| s.to_string())
.collect();
pub async fn send_dm(msg: String, args: Args, owners: Vec<u64>) {
let token: String =
if !args.test { std::env::var("ASSISTANT_TOKEN") .expect("Missing ASSISTANT_TOKEN env var!") }
else { std::env::var("ASSISTANT_TOKEN_TEST").expect("Missing ASSISTANT_TOKEN_TEST env var!") };
if !args.test { env::var("ASSISTANT_TOKEN") .expect("Missing ASSISTANT_TOKEN env var!") }
else { env::var("ASSISTANT_TOKEN_TEST").expect("Missing ASSISTANT_TOKEN_TEST env var!") };
let http = Http::new(&token);
let c_msg = CreateMessage::new().content(msg);
for uid in uids_split {
let user = UserId::new(uid.parse::<u64>().unwrap());
for uid in owners {
let user = UserId::new(uid);
let _ = user.dm(http.as_ref(), c_msg.clone()).await;
}
}
src/websocket.rs
+6 -6
View File
@@ -98,17 +98,17 @@ async fn receive_response() -> Option<Value> {
}
pub async fn start(args: Args) {
pub async fn start(args: Args, owners: Vec<u64>) {
rs_println!("Starting local websocket...");
let ip = format!("127.0.0.1:{}", args.port);
let listener = TcpListener::bind(&ip).await.unwrap();
rs_println!("WebSocket server running on ws://{}", ip);
tokio::spawn(handle_connections(listener, args));
tokio::spawn(handle_connections(listener, args, owners));
}
async fn handle_connections(listener: TcpListener, args: Args) {
async fn handle_connections(listener: TcpListener, args: Args, owners: Vec<u64>) {
while let Ok((stream, _)) = listener.accept().await {
let ws_stream = accept_async(stream).await.unwrap();
let (sender, receiver) = ws_stream.split();
@@ -120,13 +120,13 @@ async fn handle_connections(listener: TcpListener, args: Args) {
set_receiver(receiver_arc.clone()).await;
while let Some(Ok(msg)) = receiver_arc.lock().await.as_mut().unwrap().next().await {
handle_message(msg, args.clone()).await;
handle_message(msg, args.clone(), owners.clone()).await;
}
}
}
async fn handle_message(msg: tungstenite::protocol::Message, args: Args) {
async fn handle_message(msg: tungstenite::protocol::Message, args: Args, owners: Vec<u64>) {
match msg {
tungstenite::Message::Text(text) => {
rs_println!("Received from Python: {}", text);
@@ -134,7 +134,7 @@ async fn handle_message(msg: tungstenite::protocol::Message, args: Args) {
if text.starts_with("json:") {
let t_json: Value = serde_json::from_str(&text[5..]).unwrap();
if t_json.get("error").is_some() {
send_dm("Unknown internal Python error occurred!".to_string(), args).await;
send_dm("Unknown internal Python error occurred!".to_string(), args, owners).await;
}
}