use futures::stream::SplitStream; use tokio::sync::Mutex; use futures::SinkExt; use tokio::net::{TcpListener, TcpStream}; use tokio_tungstenite::WebSocketStream; use tokio_tungstenite::{accept_async, tungstenite}; use futures::StreamExt; use std::sync::Arc; use serde_json::{Value, json}; use crate::db::bot_data::Data; use crate::messages::send_dm_min; use crate::rs_println; use crate::Args; type Sender = Arc, tungstenite::Message>>>>; type Receiver = Arc>>>>; static mut GLOBAL_SENDER: Option = None; static mut GLOBAL_RECEIVER: Option = None; static mut REPLY_HELLO: bool = false; pub static mut HAS_CONNECTED: bool = false; async fn set_sender(sender: Sender) { unsafe { GLOBAL_SENDER = Some(sender); } } async fn set_receiver(receiver: Receiver) { unsafe { GLOBAL_RECEIVER = Some(receiver); } } #[allow(static_mut_refs)] pub async fn send_msg(msg: &str) { unsafe { if let Some(sender) = &GLOBAL_SENDER { let mut sender = sender.lock().await; if let Some(s) = sender.as_mut() { s.send(tungstenite::Message::Text(msg.to_string().into())).await.unwrap(); } } } } #[allow(static_mut_refs)] pub async fn send_cmd_json(func_name: &str, func_args: Option, print_output: bool) -> Option { unsafe { let Some(sender) = &GLOBAL_SENDER else { return None }; let mut sender = sender.lock().await; let s = sender.as_mut()?; let unw_args = func_args.unwrap_or(json!([])); let json_str = format!( "json:{{\"type\": \"function\", \"value\":\"{}\", \"args\": {}, \"print\": {print_output}}}", func_name, unw_args ); if s.send(tungstenite::Message::Text(json_str.into())).await.is_err() { return None; } let r = receive_response().await; if let Some(rs) = r.clone() { if !rs.get("print").unwrap_or(&json![false]).as_bool().unwrap() { return r; } } if ::parse().dev { rs_println!("Received from Python: [RESPONSE] {:?}", r); } if r.is_none() { rs_println!("[IMPORTANT] Response from Python is None!"); } return r; } } #[allow(static_mut_refs)] async fn receive_response() -> Option { unsafe { let Some(receiver) = &GLOBAL_RECEIVER else { return None }; let mut receiver = receiver.lock().await; let r = receiver.as_mut()?; let Some(Ok(msg)) = r.next().await else { return None }; let tungstenite::Message::Text(response) = msg else { return None }; if let Some(stripped) = response.strip_prefix("json:") { return serde_json::from_str(stripped).ok(); } else { return serde_json::from_str(&response).ok(); } } } pub async fn start(data: &Data) { rs_println!("Running local websocket..."); let ip = format!("127.0.0.1:{}", data.args.port); let listener = TcpListener::bind(&ip).await.unwrap(); rs_println!("WebSocket server running on ws://{}", ip); tokio::spawn(handle_connections( listener, data.env_vars.token.clone(), data.env_vars.bot_owners.clone() )); } async fn handle_connections( listener: TcpListener, token: String, bot_owners: Vec ) { while let Ok((stream, _)) = listener.accept().await { let ws_stream = accept_async(stream).await.unwrap(); let (sender, receiver) = ws_stream.split(); let sender_arc = Arc::new(Mutex::new(Some(sender))); let receiver_arc = Arc::new(Mutex::new(Some(receiver))); set_sender(sender_arc.clone()).await; set_receiver(receiver_arc.clone()).await; while let Some(Ok(msg)) = receiver_arc.lock().await.as_mut().unwrap().next().await { handle_message(msg, token.clone(), bot_owners.clone()).await; } } } async fn handle_message( msg: tungstenite::protocol::Message, token: String, bot_owners: Vec, ) { match msg { tungstenite::Message::Text(text) => { rs_println!("Received from Python: {}", text); if let Some(stripped) = text.strip_prefix("json:") { let t_json: Value = serde_json::from_str(stripped).unwrap(); if t_json.get("error").is_some() { send_dm_min( "Unknown internal Python error occurred: Websocket response error".to_string(), token, bot_owners ).await; } } unsafe { if !REPLY_HELLO { send_msg("[Connection test] Hello from Rust!").await; REPLY_HELLO = true; HAS_CONNECTED = true; } } }, _ => {} } }