diff --git a/src/main.rs b/src/main.rs index 1b9ed36..5094c33 100644 --- a/src/main.rs +++ b/src/main.rs @@ -6,8 +6,6 @@ use tokio::sync::mpsc; use serde::Deserialize; use std::fs; -use rand::{thread_rng, Rng}; - use std::sync::atomic::{AtomicBool, Ordering}; use colored::*; use tokio_socks::tcp::Socks5Stream; @@ -29,7 +27,7 @@ struct Config { // Proxy use_proxy: bool, - proxy_type: Option, + // proxy_type: Option, proxy_addr: Option, proxy_port: Option, proxy_username: Option, @@ -40,16 +38,23 @@ struct Config { } mod mods { + pub mod proxy; + pub mod tls; + pub mod handler; pub mod sasl; pub mod sed; pub mod ascii; pub mod vomit; +// pub mod invade; } +use mods::proxy::proxy_exec; +use mods::tls::tls_exec; +use mods::handler::handler; use mods::sasl::{start_sasl_auth, handle_sasl_messages}; use mods::sed::{SedCommand, MessageBuffer}; use mods::ascii::handle_ascii_command; -use mods::vomit::{handle_vomit_command}; - +use mods::vomit::handle_vomit_command; +//use mods::invade::{handle_invade_command}; #[tokio::main(flavor = "multi_thread", worker_threads = 12)] async fn main() -> Result<(), Box> { @@ -121,58 +126,6 @@ fn loaded_config() -> Result> { Ok(config) } -/// Establish a TLS connection to the server -async fn tls_exec(config: &Config, tcp_stream: TcpStream) -> Result, Box> { - let tls_builder = NTlsConnector::builder().danger_accept_invalid_certs(true).build().unwrap(); - let tls_connector = TlsConnector::from(tls_builder); - Ok(tls_connector.connect(&config.server, tcp_stream).await.unwrap()) - -} - -/// Establish a connection to the proxy -async fn proxy_exec(config: &Config) -> Result> { - let proxy_addr = match config.proxy_addr.as_ref() { - Some(addr) => addr, - None => "127.0.0.1", - }; - let proxy_port = config.proxy_port.unwrap_or(9050); - let proxy = format!("{}:{}", proxy_addr, proxy_port); - let server = format!("{}:{}", config.server, config.port); - let proxy_stream = TcpStream::connect(proxy).await.unwrap(); - let username = config.proxy_username.clone().unwrap(); - let password = config.proxy_password.clone().unwrap(); - let tcp_stream = if !&username.is_empty() && !password.is_empty() { - let tcp_stream = Socks5Stream::connect_with_password_and_socket(proxy_stream, server, &username, &password).await.unwrap(); - tcp_stream - } else { - let tcp_stream = Socks5Stream::connect_with_socket(proxy_stream, server).await.unwrap(); - tcp_stream - }; - let tcp_stream = tcp_stream.into_inner(); - - Ok(tcp_stream) -} - -/// Handle the connection to the server -async fn handler(stream: S, config: Config) -> Result<(), Box> where S: AsyncRead + AsyncWrite + Unpin + Send + 'static { - let (reader, writer) = split(stream); - let (tx, rx) = mpsc::channel(1000); - - let read_task = tokio::spawn(async move { - readmsg(reader, tx).await; - }); - - let message_buffer = MessageBuffer::new(1000); - - let write_task = tokio::spawn(async move { - writemsg(writer, rx, &config, message_buffer).await; - }); - - //let _ = tokio::try_join!(read_task, write_task); - tokio::try_join!(read_task, write_task).map_err(|e| Box::new(e) as Box)?; - Ok(()) -} - /// Read messages from the server async fn readmsg(mut reader: tokio::io::ReadHalf, tx: tokio::sync::mpsc::Sender) where S: AsyncRead + Unpin { let mut buf = vec![0; 4096]; @@ -288,10 +241,15 @@ async fn writemsg(mut writer: tokio::io::WriteHalf, mut rx: tokio::sync::m let _ = handle_ascii_command(&mut writer, config, &msg_content, channel).await; } - + // vomit if msg_content.starts_with("%vomit") { let _ = handle_vomit_command(&mut writer, config, &msg_content, channel).await; } + + // invade +// if msg_content.starts_with("%invade") { +// let _ = handle_vomit_command(&mut writer, config, &msg_content, channel).await; +// } // other commands here } } diff --git a/src/mods/handler.rs b/src/mods/handler.rs new file mode 100644 index 0000000..4edab29 --- /dev/null +++ b/src/mods/handler.rs @@ -0,0 +1,24 @@ +// mods/handler.rs +use tokio::io::{AsyncRead, AsyncWrite, split}; +use tokio::sync::mpsc; +use crate::{Config, readmsg, writemsg, MessageBuffer}; + +/// Handle the connection to the server +pub async fn handler(stream: S, config: Config) -> Result<(), Box> where S: AsyncRead + AsyncWrite + Unpin + Send + 'static { + let (reader, writer) = split(stream); + let (tx, rx) = mpsc::channel(1000); + + let read_task = tokio::spawn(async move { + readmsg(reader, tx).await; + }); + + let message_buffer = MessageBuffer::new(1000); + + let write_task = tokio::spawn(async move { + writemsg(writer, rx, &config, message_buffer).await; + }); + + //let _ = tokio::try_join!(read_task, write_task); + tokio::try_join!(read_task, write_task).map_err(|e| Box::new(e) as Box)?; + Ok(()) +} diff --git a/src/mods/invade.rs b/src/mods/invade.rs new file mode 100644 index 0000000..85d7028 --- /dev/null +++ b/src/mods/invade.rs @@ -0,0 +1,5 @@ +// mods/invade.rs + + + + diff --git a/src/mods/proxy.rs b/src/mods/proxy.rs new file mode 100644 index 0000000..ca49aa0 --- /dev/null +++ b/src/mods/proxy.rs @@ -0,0 +1,28 @@ +// mods/proxy.rs +use tokio::net::TcpStream; +use tokio_socks::tcp::Socks5Stream; +use crate::Config; + +/// Establish a connection to the proxy +pub async fn proxy_exec(config: &Config) -> Result> { + let proxy_addr = match config.proxy_addr.as_ref() { + Some(addr) => addr, + None => "127.0.0.1", + }; + let proxy_port = config.proxy_port.unwrap_or(9050); + let proxy = format!("{}:{}", proxy_addr, proxy_port); + let server = format!("{}:{}", config.server, config.port); + let proxy_stream = TcpStream::connect(proxy).await.unwrap(); + let username = config.proxy_username.clone().unwrap(); + let password = config.proxy_password.clone().unwrap(); + let tcp_stream = if !&username.is_empty() && !password.is_empty() { + let tcp_stream = Socks5Stream::connect_with_password_and_socket(proxy_stream, server, &username, &password).await.unwrap(); + tcp_stream + } else { + let tcp_stream = Socks5Stream::connect_with_socket(proxy_stream, server).await.unwrap(); + tcp_stream + }; + let tcp_stream = tcp_stream.into_inner(); + + Ok(tcp_stream) +} diff --git a/src/mods/tls.rs b/src/mods/tls.rs new file mode 100644 index 0000000..ddec428 --- /dev/null +++ b/src/mods/tls.rs @@ -0,0 +1,12 @@ +// mods/tls.rs +use tokio::net::TcpStream; +use tokio_native_tls::{TlsConnector, native_tls::TlsConnector as NTlsConnector}; +use crate::Config; + +// Establish a TLS connection to the server +pub async fn tls_exec(config: &Config, tcp_stream: TcpStream) -> Result, Box> { + let tls_builder = NTlsConnector::builder().danger_accept_invalid_certs(true).build().unwrap(); + let tls_connector = TlsConnector::from(tls_builder); + Ok(tls_connector.connect(&config.server, tcp_stream).await.unwrap()) + +}