diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..6313b56 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +* text=auto eol=lf diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 0000000..cd13d94 --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1,2 @@ +# Tabs are optimal +hard_tabs = true diff --git a/src/colors.rs b/src/colors.rs index f363ed1..dc6b7a3 100644 --- a/src/colors.rs +++ b/src/colors.rs @@ -1,4 +1,4 @@ pub const GREEN: &str = "\x1b[0;32m"; pub const YELLOW: &str = "\x1b[0;33m"; pub const RED: &str = "\x1b[0;31m"; -pub const RESET: &str = "\x1b[0m"; \ No newline at end of file +pub const RESET: &str = "\x1b[0m"; diff --git a/src/config.rs b/src/config.rs index 0628b90..4ed2f52 100644 --- a/src/config.rs +++ b/src/config.rs @@ -5,68 +5,68 @@ use std::io::{ErrorKind, Read}; #[derive(Deserialize)] pub struct Config { - pub database: Database, - pub rescanner: Rescanner + pub database: Database, + pub rescanner: Rescanner, } #[derive(Deserialize)] pub struct Database { - pub url: String, - pub port: u16, - pub table: String, - pub user: String, - pub password: String, + pub url: String, + pub port: u16, + pub table: String, + pub user: String, + pub password: String, } #[derive(Deserialize)] pub struct Rescanner { - pub repeat: bool, - pub rescan_delay: u64, - pub port_range_start: u16, - pub port_range_end: u16, + pub repeat: bool, + pub rescan_delay: u64, + pub port_range_start: u16, + pub port_range_end: u16, } impl Default for Rescanner { - fn default() -> Self { - Self { - repeat: true, - rescan_delay: 60, - port_range_start: 25560, - port_range_end: 25570, - } - } + fn default() -> Self { + Self { + repeat: true, + rescan_delay: 60, + port_range_start: 25560, + port_range_end: 25570, + } + } } impl Rescanner { - pub fn total_ports(&self) -> u16 { - let start = self.port_range_start; - let end = self.port_range_end; + pub fn total_ports(&self) -> u16 { + let start = self.port_range_start; + let end = self.port_range_end; - if start > end { - println!("{RED}port_range_start cannot be greater than port_range_end!{RESET}"); - std::process::exit(1); - } + if start > end { + println!("{RED}port_range_start cannot be greater than port_range_end!{RESET}"); + std::process::exit(1); + } - if end - start == 0 { - return 1; - } + if end - start == 0 { + return 1; + } - end - start - } + end - start + } } pub fn load_config(path: String) -> Config { - let mut file = match File::open(&path) { - Ok(file) => file, - Err(error) => match error.kind() { - ErrorKind::NotFound => panic!("{} not found!", &path), - ErrorKind::PermissionDenied => panic!("File permissions invalid for {}!", &path), - _ => panic!("Unknown config error!"), - } - }; + let mut file = match File::open(&path) { + Ok(file) => file, + Err(error) => match error.kind() { + ErrorKind::NotFound => panic!("{} not found!", &path), + ErrorKind::PermissionDenied => panic!("File permissions invalid for {}!", &path), + _ => panic!("Unknown config error!"), + }, + }; - let mut contents = String::new(); - file.read_to_string(&mut contents).unwrap_or_default(); + let mut contents = String::new(); + file.read_to_string(&mut contents).unwrap_or_default(); - toml::from_str(&contents).expect("Failed to parse config!") -} \ No newline at end of file + toml::from_str(&contents).expect("Failed to parse config!") +} diff --git a/src/database.rs b/src/database.rs index 96b8976..4c3e0fa 100644 --- a/src/database.rs +++ b/src/database.rs @@ -1,34 +1,35 @@ -use std::rc::Rc; use crate::colors::{RED, RESET}; use crate::response::Server; use sqlx::{Error, PgPool, Pool, Postgres, Row}; +use std::rc::Rc; use std::time::{SystemTime, UNIX_EPOCH}; pub async fn connect(url: &str) -> Pool { - match PgPool::connect(url).await { - Ok(pool) => pool, - Err(e) => panic!("{RED}Unable to connect to database: {e}{RESET}"), - } + match PgPool::connect(url).await { + Ok(pool) => pool, + Err(e) => panic!("{RED}Unable to connect to database: {e}{RESET}"), + } } // TODO! Return a stream of results instead of a Vec for performance pub async fn fetch_servers(pool: &Pool) -> Result, Error> { - // Sort results by oldest servers first - sqlx::query("SELECT address FROM servers ORDER BY lastseen DESC") - .fetch_all(&*pool) - .await? - .into_iter() - .map(|row| row.try_get(0) ) - .collect() + // Sort results by oldest servers first + sqlx::query("SELECT address FROM servers ORDER BY lastseen DESC") + .fetch_all(&*pool) + .await? + .into_iter() + .map(|row| row.try_get(0)) + .collect() } pub async fn update(server: Server, conn: &Rc) -> anyhow::Result<()> { - let timestamp = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs() as i32; - let mut transaction = conn.begin().await?; - let address: &str = server.address.as_str(); + let timestamp = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs() as i32; + let mut transaction = conn.begin().await?; + let address: &str = server.address.as_str(); - // Update server - sqlx::query("UPDATE servers SET + // Update server + sqlx::query( + "UPDATE servers SET version = $1, protocol = $2, icon = $3, @@ -39,24 +40,25 @@ pub async fn update(server: Server, conn: &Rc) -> anyhow::Result<()> { onlineplayers = $8, maxplayers = $9 WHERE address = $10 - AND port = $11") - .bind(server.version) - .bind(server.protocol,) - .bind(server.icon) - .bind(server.motd) - .bind(server.prevents_reports) - .bind(server.enforces_secure_chat) - .bind(timestamp) - .bind(server.online_players) - .bind(server.max_players) - .bind(address) - .bind(server.port) - .execute(&mut *transaction) - .await?; + AND port = $11", + ) + .bind(server.version) + .bind(server.protocol) + .bind(server.icon) + .bind(server.motd) + .bind(server.prevents_reports) + .bind(server.enforces_secure_chat) + .bind(timestamp) + .bind(server.online_players) + .bind(server.max_players) + .bind(address) + .bind(server.port) + .execute(&mut *transaction) + .await?; - // Upsert players - for player in server.players { - sqlx::query("INSERT INTO playerhistory (address, port, playeruuid, playername, firstseen, lastseen) VALUES ($1, $2, $3, $4, $5, $6) + // Upsert players + for player in server.players { + sqlx::query("INSERT INTO playerhistory (address, port, playeruuid, playername, firstseen, lastseen) VALUES ($1, $2, $3, $4, $5, $6) ON CONFLICT (address, port, playeruuid) DO UPDATE SET lastseen = EXCLUDED.lastseen, playername = EXCLUDED.playername") @@ -68,11 +70,11 @@ pub async fn update(server: Server, conn: &Rc) -> anyhow::Result<()> { .bind(timestamp) .execute(&mut *transaction) .await?; - }; + } - // Update mods - for mods in server.mods { - sqlx::query("INSERT INTO mods (address, port, modid, modmarker) VALUES ($1, $2, $3, $4) ON CONFLICT (address, port, modid) DO NOTHING") + // Update mods + for mods in server.mods { + sqlx::query("INSERT INTO mods (address, port, modid, modmarker) VALUES ($1, $2, $3, $4) ON CONFLICT (address, port, modid) DO NOTHING") .bind(address) .bind(server.port) .bind(mods.mod_id) @@ -81,7 +83,7 @@ pub async fn update(server: Server, conn: &Rc) -> anyhow::Result<()> { .bind(timestamp) .execute(&mut *transaction) .await?; - }; - - Ok(transaction.commit().await?) -} \ No newline at end of file + } + + Ok(transaction.commit().await?) +} diff --git a/src/main.rs b/src/main.rs index 197f67f..4718262 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,136 +1,153 @@ -mod database; +mod colors; mod config; +mod database; mod ping; mod response; -mod colors; -use std::rc::Rc; use colors::{GREEN, RED, RESET, YELLOW}; use config::load_config; use database::{connect, fetch_servers}; use indicatif::{ProgressBar, ProgressStyle}; -use std::time::Duration; use sqlx::{Pool, Postgres}; +use std::rc::Rc; +use std::time::Duration; use tokio::sync::Semaphore; #[derive(Clone)] struct State { - pool: Rc>, - semaphore: Rc, - progress_bar: Rc + pool: Rc>, + semaphore: Rc, + progress_bar: Rc, } #[tokio::main] async fn main() { - let config_file = std::env::args().nth(1).unwrap_or("config.toml".to_string()); - println!("{GREEN}[INFO] Using config file {}{RESET}", config_file); - let config = load_config(config_file); + let config_file = std::env::args().nth(1).unwrap_or("config.toml".to_string()); + println!("{GREEN}[INFO] Using config file {}{RESET}", config_file); + let config = load_config(config_file); - // Create database URL - let database_url = format!("postgresql://{}:{}@{}:{}/{}", - config.database.user, - config.database.password, - config.database.url, - config.database.port, - config.database.table); + // Create database URL + let database_url = format!( + "postgresql://{}:{}@{}:{}/{}", + config.database.user, + config.database.password, + config.database.url, + config.database.port, + config.database.table + ); - let port_start = config.rescanner.port_range_start; - let port_end = config.rescanner.port_range_end; - let total_ports = config.rescanner.total_ports(); + let port_start = config.rescanner.port_range_start; + let port_end = config.rescanner.port_range_end; + let total_ports = config.rescanner.total_ports(); - if total_ports > 10 { - println!("{RED}[WARN] Large amount of ports! Scans will take exponentially longer for each port to scan!{RESET}"); - } + if total_ports > 10 { + println!("{RED}[WARN] Large amount of ports! Scans will take exponentially longer for each port to scan!{RESET}"); + } - if !config.rescanner.repeat { - println!("{YELLOW}[WARN] Repeat is not enabled in config file! Will only scan once!{RESET}"); - } + if !config.rescanner.repeat { + println!( + "{YELLOW}[WARN] Repeat is not enabled in config file! Will only scan once!{RESET}" + ); + } - println!("{GREEN}[INFO] Scanning port range {} - {} ({} port(s) per host){RESET}", - port_start, - port_end, - total_ports); + println!( + "{GREEN}[INFO] Scanning port range {} - {} ({} port(s) per host){RESET}", + port_start, port_end, total_ports + ); - let semaphore = Rc::new(Semaphore::new(3000)); - - loop { - let pool = Rc::new(connect(database_url.as_str()).await); + let semaphore = Rc::new(Semaphore::new(3000)); - let servers = match fetch_servers(&pool).await { - Ok(servers) => { - println!("{GREEN}[INFO] Found {} servers to rescan!{RESET}", servers.len()); - servers - }, - Err(_) => { - println!("{RED}[ERROR] Failed to fetch servers! Waiting 30 seconds and retrying...{RESET}"); - tokio::time::sleep(Duration::from_secs(30)).await; - continue; - } - }; + loop { + let pool = Rc::new(connect(database_url.as_str()).await); - let style = ProgressStyle::with_template("[{elapsed}] [{bar:40.white/blue}] {pos:>7}/{len:7}").unwrap().progress_chars("=>-"); - let progress_bar = Rc::new(ProgressBar::new(servers.len() as u64).with_style(style)); + let servers = match fetch_servers(&pool).await { + Ok(servers) => { + println!( + "{GREEN}[INFO] Found {} servers to rescan!{RESET}", + servers.len() + ); + servers + } + Err(_) => { + println!("{RED}[ERROR] Failed to fetch servers! Waiting 30 seconds and retrying...{RESET}"); + tokio::time::sleep(Duration::from_secs(30)).await; + continue; + } + }; - let state = Rc::new(State { - // Pool isn't used anywhere else except for inside the futures so it's safe to move the value - pool: pool.clone(), - semaphore: semaphore.clone(), - progress_bar: progress_bar.clone(), - }); + let style = + ProgressStyle::with_template("[{elapsed}] [{bar:40.white/blue}] {pos:>7}/{len:7}") + .unwrap() + .progress_chars("=>-"); + let progress_bar = Rc::new(ProgressBar::new(servers.len() as u64).with_style(style)); - let servers = servers - .iter() - .map(|ip| (port_start..=port_end).map(|port| run((ip.to_owned(), port), Rc::clone(&state))).collect::>()) - .flatten() - .collect::>(); + let state = Rc::new(State { + // Pool isn't used anywhere else except for inside the futures so it's safe to move the value + pool: pool.clone(), + semaphore: semaphore.clone(), + progress_bar: progress_bar.clone(), + }); - let results = futures::future::join_all(servers).await; + let servers = servers + .iter() + .map(|ip| { + (port_start..=port_end) + .map(|port| run((ip.to_owned(), port), Rc::clone(&state))) + .collect::>() + }) + .flatten() + .collect::>(); - let errors = results.into_iter() - .filter_map(Result::err) - .collect::>(); - - if !errors.is_empty() { - println!("{YELLOW}[INFO] Scan returned {} errors!{RESET}", errors.len()); - } - - println!("{GREEN}[INFO] Finished pinging all servers{RESET}"); - - if !config.rescanner.repeat { - println!("{GREEN}[INFO] Exiting...{RESET}"); - std::process::exit(0); - } + let results = futures::future::join_all(servers).await; - if config.rescanner.rescan_delay > 0 { - println!("{GREEN}[INFO] Waiting {} seconds before starting another scan...{RESET}", config.rescanner.rescan_delay); - tokio::time::sleep(Duration::from_secs(config.rescanner.rescan_delay)).await; - } - } + let errors = results + .into_iter() + .filter_map(Result::err) + .collect::>(); + + if !errors.is_empty() { + println!( + "{YELLOW}[INFO] Scan returned {} errors!{RESET}", + errors.len() + ); + } + + println!("{GREEN}[INFO] Finished pinging all servers{RESET}"); + + if !config.rescanner.repeat { + println!("{GREEN}[INFO] Exiting...{RESET}"); + std::process::exit(0); + } + + if config.rescanner.rescan_delay > 0 { + println!( + "{GREEN}[INFO] Waiting {} seconds before starting another scan...{RESET}", + config.rescanner.rescan_delay + ); + tokio::time::sleep(Duration::from_secs(config.rescanner.rescan_delay)).await; + } + } } async fn run(host: (String, u16), state: Rc) -> Result<(), ErrorType> { - let permit = state.semaphore.acquire().await; + let permit = state.semaphore.acquire().await; - match ping::ping_server(&host).await { - Ok(results) => { - match response::parse_response(results, &host) { - Ok(response) => { - match database::update(response, &state.pool).await { - Ok(_) => (), - _ => return Err(ErrorType::DatabaseError) - } - } - _ => return Err(ErrorType::ParsingError) - } - } - _ => return Err(ErrorType::ConnectionRefused) - } + match ping::ping_server(&host).await { + Ok(results) => match response::parse_response(results, &host) { + Ok(response) => match database::update(response, &state.pool).await { + Ok(_) => (), + _ => return Err(ErrorType::DatabaseError), + }, + _ => return Err(ErrorType::ParsingError), + }, + _ => return Err(ErrorType::ConnectionRefused), + } - Ok(state.progress_bar.inc(1)) + Ok(state.progress_bar.inc(1)) } enum ErrorType { - ConnectionRefused, - ParsingError, - DatabaseError -} \ No newline at end of file + ConnectionRefused, + ParsingError, + DatabaseError, +} diff --git a/src/ping.rs b/src/ping.rs index 430fa4a..e642652 100644 --- a/src/ping.rs +++ b/src/ping.rs @@ -5,61 +5,62 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::TcpStream; const PAYLOAD: [u8; 9] = [ - 6, // Size: Amount of bytes in the message - 0, // ID: Has to be 0 - 0, // Protocol Version: Can be anything as long as it's a valid varint - 0, // Server address - 0, 0, // Port: Can be anything (Notchian servers don't this) - 1, // Next state: 1 for status, 2 for login. Therefore, has to be 1 - 1, // Size - 0, // ID + 6, // Size: Amount of bytes in the message + 0, // ID: Has to be 0 + 0, // Protocol Version: Can be anything as long as it's a valid varint + 0, // Server address + 0, 0, // Port: Can be anything (Notchian servers don't this) + 1, // Next state: 1 for status, 2 for login. Therefore, has to be 1 + 1, // Size + 0, // ID ]; pub async fn ping_server(host: &(String, u16)) -> anyhow::Result { - let address = format!("{}:{}", host.0, host.1); - let socket = SocketAddr::from_str(address.as_str())?; + let address = format!("{}:{}", host.0, host.1); + let socket = SocketAddr::from_str(address.as_str())?; - // Connect and create buffer - let mut stream = tokio::time::timeout(Duration::from_secs(3), TcpStream::connect(&socket)).await??; - let mut buffer = [0; 1024]; + // Connect and create buffer + let mut stream = + tokio::time::timeout(Duration::from_secs(3), TcpStream::connect(&socket)).await??; + let mut buffer = [0; 1024]; - // Send payload - stream.write(&PAYLOAD).await?; - let mut total_read = stream.read(&mut buffer).await?; - - // Decode - let (varint, length) = decode(&buffer); - let bytes_needed = varint + length as usize; - let mut output = vec![]; - output.extend_from_slice(&buffer[..total_read]); - let json = decode(&(buffer[(length + 1).into()..])); - - // Read everything - while total_read < bytes_needed { - let read = stream.read(&mut buffer).await?; - output.extend_from_slice(&buffer[..read]); - total_read += read; - } + // Send payload + stream.write(&PAYLOAD).await?; + let mut total_read = stream.read(&mut buffer).await?; - // Explicitly shutdown stream - stream.shutdown().await?; - - Ok(String::from_utf8_lossy(&output[(length + 1 + json.1).into()..]).to_string()) + // Decode + let (varint, length) = decode(&buffer); + let bytes_needed = varint + length as usize; + let mut output = vec![]; + output.extend_from_slice(&buffer[..total_read]); + let json = decode(&(buffer[(length + 1).into()..])); + + // Read everything + while total_read < bytes_needed { + let read = stream.read(&mut buffer).await?; + output.extend_from_slice(&buffer[..read]); + total_read += read; + } + + // Explicitly shutdown stream + stream.shutdown().await?; + + Ok(String::from_utf8_lossy(&output[(length + 1 + json.1).into()..]).to_string()) } fn decode(bytes: &[u8]) -> (usize, u8) { - let mut val: usize = 0; - let mut count: u8 = 0; + let mut val: usize = 0; + let mut count: u8 = 0; - for b in bytes { - val |= ((b & 0x7f) as usize) << count; + for b in bytes { + val |= ((b & 0x7f) as usize) << count; - if (b >> 7) != 1 { - break; - } + if (b >> 7) != 1 { + break; + } - count += 7; - } + count += 7; + } - (val, (count / 7) + 1) -} \ No newline at end of file + (val, (count / 7) + 1) +} diff --git a/src/response.rs b/src/response.rs index 4bbcfee..92c7d21 100644 --- a/src/response.rs +++ b/src/response.rs @@ -4,101 +4,101 @@ use std::str::FromStr; #[allow(dead_code)] #[derive(Debug)] pub struct Server { - pub address: String, - pub port: i32, - pub version: Option, - pub protocol: Option, - pub motd: Option, - pub icon: Option, - pub prevents_reports: Option, - pub enforces_secure_chat: Option, - pub online_players: Option, - pub max_players: Option, - pub software: Option, - pub mods: Vec, - pub players: Vec, + pub address: String, + pub port: i32, + pub version: Option, + pub protocol: Option, + pub motd: Option, + pub icon: Option, + pub prevents_reports: Option, + pub enforces_secure_chat: Option, + pub online_players: Option, + pub max_players: Option, + pub software: Option, + pub mods: Vec, + pub players: Vec, } #[allow(dead_code)] #[derive(Debug)] pub struct Player { - pub username: Option, - // TODO! Replace this with a UUID type - pub uuid: Option + pub username: Option, + // TODO! Replace this with a UUID type + pub uuid: Option, } #[allow(dead_code)] #[derive(Debug)] pub struct Mod { - pub mod_id: Option, - pub mod_name: Option, + pub mod_id: Option, + pub mod_name: Option, } pub fn parse_response(response: String, host: &(String, u16)) -> anyhow::Result { - let json = Value::from_str(response.as_str())?; + let json = Value::from_str(response.as_str())?; - let mut version: Option = None; - let mut protocol: Option = None; - let mut online_players: Option = None; - let mut max_players: Option = None; - let mut players = vec![]; - let mut mods = vec![]; + let mut version: Option = None; + let mut protocol: Option = None; + let mut online_players: Option = None; + let mut max_players: Option = None; + let mut players = vec![]; + let mut mods = vec![]; - let icon = json["icon"].as_str().map(String::from); - let prevents_reports = json["preventsChatReports"].as_bool(); - let enforces_secure_chat = json["enforcesChat"].as_bool(); + let icon = json["icon"].as_str().map(String::from); + let prevents_reports = json["preventsChatReports"].as_bool(); + let enforces_secure_chat = json["enforcesChat"].as_bool(); - // Version object - if let Some(value) = json.get("version") { - version = value["name"].as_str().map(String::from); + // Version object + if let Some(value) = json.get("version") { + version = value["name"].as_str().map(String::from); - if let Some(p) = value["protocol"].as_i64() { - protocol = Some(p as i32) - } - } + if let Some(p) = value["protocol"].as_i64() { + protocol = Some(p as i32) + } + } - // Players object - if let Some(value) = json.get("players") { - if let Some(online) = value["online"].as_i64() { - online_players = Some(online as i32); - } + // Players object + if let Some(value) = json.get("players") { + if let Some(online) = value["online"].as_i64() { + online_players = Some(online as i32); + } - if let Some(max) = value["max"].as_i64() { - max_players = Some(max as i32) - } + if let Some(max) = value["max"].as_i64() { + max_players = Some(max as i32) + } - if let Some(player_sample) = value["sample"].as_array() { - player_sample.iter().for_each(|player| { - players.push(Player { - username: player["name"].as_str().map(String::from), - uuid: player["id"].as_str().map(String::from) - }) - }) - } - } + if let Some(player_sample) = value["sample"].as_array() { + player_sample.iter().for_each(|player| { + players.push(Player { + username: player["name"].as_str().map(String::from), + uuid: player["id"].as_str().map(String::from), + }) + }) + } + } - if let Some(mods_array) = json["mods"].as_array() { - mods_array.iter().for_each(|m| { - mods.push(Mod { - mod_name: m["modid"].as_str().map(String::from), - mod_id: m["version"].as_str().map(String::from) - }) - }) - } + if let Some(mods_array) = json["mods"].as_array() { + mods_array.iter().for_each(|m| { + mods.push(Mod { + mod_name: m["modid"].as_str().map(String::from), + mod_id: m["version"].as_str().map(String::from), + }) + }) + } - Ok(Server { - address: host.0.to_string(), - port: host.1 as i32, - version, - protocol, - motd: None, - icon, - prevents_reports, - enforces_secure_chat, - online_players, - max_players, - software: None, - mods, - players - }) -} \ No newline at end of file + Ok(Server { + address: host.0.to_string(), + port: host.1 as i32, + version, + protocol, + motd: None, + icon, + prevents_reports, + enforces_secure_chat, + online_players, + max_players, + software: None, + mods, + players, + }) +}