From 0543cabce26210bc8eb50bc4e3012a93f570c15f Mon Sep 17 00:00:00 2001 From: Cain <75994858+cainthebest@users.noreply.github.com> Date: Sat, 13 Jan 2024 09:18:56 +0000 Subject: [PATCH] refactor: clean up --- crates/cli/src/main.rs | 140 ++++++++++++++++++++----------- crates/lib/src/capture/mod.rs | 24 +++--- crates/lib/src/capture/packet.rs | 94 ++++----------------- crates/lib/src/capture/pcap.rs | 23 ++--- crates/lib/src/capture/writer.rs | 18 ++-- crates/lib/src/socket.rs | 9 +- 6 files changed, 137 insertions(+), 171 deletions(-) diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 969e5e3..df7078a 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -1,6 +1,9 @@ -use std::net::{IpAddr, ToSocketAddrs}; +use std::{ + net::{IpAddr, ToSocketAddrs}, + path::PathBuf, +}; -use clap::{Parser, ValueEnum}; +use clap::{Parser, Subcommand, ValueEnum}; use gamedig::{ games::*, protocols::types::{CommonResponse, ExtraRequestSettings, TimeoutSettings}, @@ -13,42 +16,68 @@ use self::error::{Error, Result}; // NOTE: For some reason without setting long_about here the doc comment for // ExtraRequestSettings gets set as the about for the CLI. #[derive(Debug, Parser)] -#[command(author, version, about, long_about = None)] +#[command(author, version, long_about = None, about = r" + + _____ _____ _ _____ _ _____ + / ____| | __ \(_) / ____| | |_ _| +| | __ __ _ _ __ ___ ___| | | |_ __ _ | | | | | | +| | |_ |/ _` | '_ ` _ \ / _ \ | | | |/ _` | | | | | | | +| |__| | (_| | | | | | | __/ |__| | | (_| | | |____| |____ _| |_ + \_____|\__,_|_| |_| |_|\___|_____/|_|\__, | \_____|______|_____| + __/ | + |___/ + + A command line interface for querying game servers. + Copyright (C) 2022 - Present GameDig Organization & Contributors + Licensed under the MIT license +")] struct Cli { - /// Unique identifier of the game for which server information is being - /// queried. - #[arg(short, long)] - game: String, + #[command(subcommand)] + action: Action, +} - /// Hostname or IP address of the server. - #[arg(short, long)] - ip: String, +#[derive(Subcommand, Debug)] +enum Action { + /// Query game server information + Query { + /// Unique identifier of the game for which server information is being + /// queried. + #[arg(short, long)] + game: String, - /// Optional query port number for the server. If not provided the default - /// port for the game is used. - #[arg(short, long)] - port: Option, + /// Hostname or IP address of the server. + #[arg(short, long)] + ip: String, - /// Flag indicating if the output should be in JSON format. - #[cfg(feature = "json")] - #[arg(short, long)] - json: bool, + /// Optional query port number for the server. If not provided the default + /// port for the game is used. + #[arg(short, long)] + port: Option, - /// Which response variant to use when outputting. - #[arg(short, long, default_value = "generic")] - output_mode: OutputMode, + /// Flag indicating if the output should be in JSON format for programmatic use. + #[cfg(feature = "json")] + #[arg(short, long)] + json: bool, - #[cfg(feature = "packet_capture")] - #[arg(short, long)] - capture: Option, + /// Which response variant to use when outputting + #[arg(short, long, default_value = "generic")] + output_mode: OutputMode, - /// Optional timeout settings for the server query. - #[command(flatten, next_help_heading = "Timeouts")] - timeout_settings: Option, + /// Optional file path for packet capture file writer + #[cfg(feature = "packet_capture")] + #[arg(short, long)] + capture: Option, - /// Optional extra settings for the server query. - #[command(flatten, next_help_heading = "Query options")] - extra_options: Option, + /// Optional timeout settings for the server query + #[command(flatten, next_help_heading = "Timeouts")] + timeout_settings: Option, + + /// Optional extra settings for the server query + #[command(flatten, next_help_heading = "Query options")] + extra_options: Option, + }, + /// Display the MIT License information + License, } #[derive(Clone, Debug, PartialEq, Eq, ValueEnum)] @@ -139,12 +168,12 @@ fn set_hostname_if_missing(host: &str, extra_options: &mut Option output_result_json(result.as_json()), + OutputMode::Generic if json => output_result_json(result.as_json()), #[cfg(feature = "json")] - OutputMode::ProtocolSpecific if args.json => output_result_json(result.as_original()), + OutputMode::ProtocolSpecific if json => output_result_json(result.as_original()), OutputMode::Generic => output_result_debug(result.as_json()), OutputMode::ProtocolSpecific => output_result_debug(result.as_original()), @@ -169,27 +198,36 @@ fn output_result_json(result: R) { } fn main() -> Result<()> { - // Parse the command line arguments let args = Cli::parse(); - // Retrieve the game based on the provided ID - let game = find_game(&args.game)?; + match args.action { + Action::Query { + game, + ip, + port, + json, + output_mode, + capture, + timeout_settings, + extra_options, + } => { + // Process the query command + let game = find_game(&game)?; + let mut extra_options = extra_options; + let ip = resolve_ip_or_domain(&ip, &mut extra_options)?; - // Extract extra options for use in setup - let mut extra_options = args.extra_options.clone(); + #[cfg(feature = "packet_capture")] + gamedig::capture::setup_capture(capture); - // Resolve the IP address - let ip = resolve_ip_or_domain(&args.ip, &mut extra_options)?; - - #[cfg(feature = "packet_capture")] - gamedig::capture::setup_capture(args.capture.clone()); - - // Query the server using game definition, parsed IP, and user command line - // flags. - let result = query_with_timeout_and_extra_settings(game, &ip, args.port, args.timeout_settings, extra_options)?; - - // Output the result in the specified format - output_result(&args, result.as_ref()); + let result = query_with_timeout_and_extra_settings(game, &ip, port, timeout_settings, extra_options)?; + output_result(output_mode, json, result.as_ref()); + } + Action::License => { + // Bake the license into the binary + // so we don't have to ship it separately + println!("{}", include_str!("../../../LICENSE.md")); + } + } Ok(()) } diff --git a/crates/lib/src/capture/mod.rs b/crates/lib/src/capture/mod.rs index 6cfd8ad..13bf7af 100644 --- a/crates/lib/src/capture/mod.rs +++ b/crates/lib/src/capture/mod.rs @@ -1,25 +1,24 @@ pub(crate) mod packet; mod pcap; -pub mod writer; +pub(crate) mod writer; -use pcap_file::pcapng::PcapNgBlock; -use writer::Writer; +use self::{pcap::Pcap, writer::Writer}; +use pcap_file::pcapng::{blocks::interface_description::InterfaceDescriptionBlock, PcapNgBlock, PcapNgWriter}; +use std::path::PathBuf; -use self::pcap::Pcap; - -pub fn setup_capture(file_name: Option) { - if let Some(file_name) = file_name { +pub fn setup_capture(file_path: Option) { + if let Some(file_path) = file_path { let file = std::fs::OpenOptions::new() .create_new(true) .write(true) - .open(file_name) + .open(file_path.with_extension("pcap")) .unwrap(); - let mut pcap_writer = pcap_file::pcapng::PcapNgWriter::new(file).unwrap(); + let mut pcap_writer = PcapNgWriter::new(file).unwrap(); // Write headers - let _ = pcap_writer.write_block( - &pcap_file::pcapng::blocks::interface_description::InterfaceDescriptionBlock { + let _ = pcap_writer.write_block( + &InterfaceDescriptionBlock { linktype: pcap_file::DataLink::ETHERNET, snaplen: 0xFFFF, options: vec![], @@ -30,6 +29,7 @@ pub fn setup_capture(file_name: Option) { let writer = Box::new(Pcap::new(pcap_writer)); attach(writer) } else { + // If no file path is provided // Do nothing } } @@ -37,7 +37,7 @@ pub fn setup_capture(file_name: Option) { /// Attaches a writer to the capture module. /// /// # Errors -/// Returns an `io::Error` if the writer is already set. +/// Returns an Error if the writer is already set. fn attach(writer: Box) { crate::socket::capture::set_writer(writer); } diff --git a/crates/lib/src/capture/packet.rs b/crates/lib/src/capture/packet.rs index 8ad2917..83910d0 100644 --- a/crates/lib/src/capture/packet.rs +++ b/crates/lib/src/capture/packet.rs @@ -30,7 +30,7 @@ pub(crate) enum Protocol { } /// Trait for handling different types of IP addresses (IPv4, IPv6). -pub trait IpAddress: Sized { +pub(crate) trait IpAddress: Sized { /// Creates an instance from a standard `IpAddr`, returning `None` if the types are incompatible. fn from_std(ip: IpAddr) -> Option; } @@ -53,8 +53,7 @@ impl CapturePacket<'_> { /// /// Returns: /// - (u16, u16): Tuple of (source port, destination port). - #[allow(dead_code)] - pub(crate) const fn ports_by_direction(&self) -> (u16, u16) { + pub(super) const fn ports_by_direction(&self) -> (u16, u16) { let (local, remote) = (self.local_address.port(), self.remote_address.port()); self.direction.order(local, remote) } @@ -63,30 +62,18 @@ impl CapturePacket<'_> { /// /// Returns: /// - (IpAddr, IpAddr): Tuple of (local IP, remote IP). - #[allow(dead_code)] - pub(crate) fn ip_addr(&self) -> (IpAddr, IpAddr) { + pub(super) fn ip_addr(&self) -> (IpAddr, IpAddr) { let (local, remote) = (self.local_address.ip(), self.remote_address.ip()); (local, remote) } - /// Retrieves IP addresses based on the packet's direction. - /// - /// Returns: - /// - (IpAddr, IpAddr): Tuple of (source IP, destination IP). - #[allow(dead_code)] - pub(crate) fn ip_addr_by_direction(&self) -> (IpAddr, IpAddr) { - let (local, remote) = self.ip_addr(); - self.direction.order(local, remote) - } - /// Retrieves IP addresses of a specific type (IPv4 or IPv6) based on the packet's direction. /// /// Panics if the IP type of the addresses does not match the requested type. /// /// Returns: /// - (T, T): Tuple of (source IP, destination IP) of the specified type in order. - #[allow(dead_code)] - pub(crate) fn ipvt_by_direction(&self) -> (T, T) { + pub(super) fn ipvt_by_direction(&self) -> (T, T) { let (local, remote) = ( T::from_std(self.local_address.ip()).expect("Incorrect IP type for local address"), T::from_std(self.remote_address.ip()).expect("Incorrect IP type for remote address"), @@ -101,7 +88,7 @@ impl Direction { /// /// Returns: /// - (T, T): Ordered tuple (source, destination). - pub(crate) const fn order(&self, source: T, remote: T) -> (T, T) { + pub(self) const fn order(&self, source: T, remote: T) -> (T, T) { match self { Direction::Send => (source, remote), Direction::Receive => (remote, source), @@ -162,65 +149,33 @@ mod tests { } #[test] - fn test_ip_addr_by_direction_ipv4() { + fn test_ip_addr() { let packet_send = CapturePacket { direction: Direction::Send, - protocol: Protocol::UDP, - local_address: &socket_addr("10.0.0.1:3000"), - remote_address: &socket_addr("10.0.0.2:3001"), + protocol: Protocol::TCP, + local_address: &socket_addr("127.0.0.1:8080"), + remote_address: &socket_addr("192.168.1.1:80"), }; let packet_receive = CapturePacket { direction: Direction::Receive, - protocol: Protocol::UDP, - local_address: &socket_addr("10.0.0.1:3000"), - remote_address: &socket_addr("10.0.0.2:3001"), + protocol: Protocol::TCP, + local_address: &socket_addr("127.0.0.1:8080"), + remote_address: &socket_addr("192.168.1.1:80"), }; assert_eq!( - packet_send.ip_addr_by_direction(), + packet_send.ip_addr(), ( - IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)), - IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2)) + IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), + IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)) ) ); assert_eq!( - packet_receive.ip_addr_by_direction(), + packet_receive.ip_addr(), ( - IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2)), - IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1)) - ) - ); - } - - #[test] - fn test_ip_addr_by_direction_ipv6() { - let packet_send = CapturePacket { - direction: Direction::Send, - protocol: Protocol::UDP, - local_address: &socket_addr("[::1]:3000"), - remote_address: &socket_addr("[::2]:3001"), - }; - - let packet_receive = CapturePacket { - direction: Direction::Receive, - protocol: Protocol::UDP, - local_address: &socket_addr("[::1]:3000"), - remote_address: &socket_addr("[::2]:3001"), - }; - - assert_eq!( - packet_send.ip_addr_by_direction(), - ( - IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)), - IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 2)) - ) - ); - assert_eq!( - packet_receive.ip_addr_by_direction(), - ( - IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 2)), - IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)) + IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), + IpAddr::V4(Ipv4Addr::new(192, 168, 1, 1)) ) ); } @@ -242,17 +197,4 @@ mod tests { std::panic::catch_unwind(|| packet.ipvt_by_direction::()); assert!(ipv6_result.is_err()); } - - #[test] - #[should_panic(expected = "Local and remote IP addresses must be of the same version")] - fn test_mismatched_ip_version_panic() { - let packet = CapturePacket { - direction: Direction::Send, - protocol: Protocol::UDP, - local_address: &socket_addr("127.0.0.1:8080"), // IPv4 - remote_address: &socket_addr("[::1]:80"), // IPv6 - }; - - packet.ip_addr_by_direction(); - } } diff --git a/crates/lib/src/capture/pcap.rs b/crates/lib/src/capture/pcap.rs index be65968..e6cd5b0 100644 --- a/crates/lib/src/capture/pcap.rs +++ b/crates/lib/src/capture/pcap.rs @@ -15,7 +15,6 @@ use super::packet::{ PACKET_SIZE, }; - const BUFFER_SIZE: usize = PACKET_SIZE - (if HEADER_SIZE_IP4 > HEADER_SIZE_IP6 { HEADER_SIZE_IP4 @@ -48,10 +47,7 @@ impl Pcap { pub(crate) fn write_transport_packet(&mut self, info: &CapturePacket, payload: &[u8]) { let mut buf = vec![0; BUFFER_SIZE]; - let (source_port, dest_port) = match info.direction { - Direction::Send => (info.local_address.port(), info.remote_address.port()), - Direction::Receive => (info.remote_address.port(), info.local_address.port()), - }; + let (source_port, dest_port) = info.ports_by_direction(); match info.protocol { Protocol::TCP => { @@ -150,13 +146,9 @@ impl Pcap { protocol: IpNextHeaderProtocol, payload: &[u8], ) -> (usize, EtherType) { - match (info.local_address.ip(), info.remote_address.ip()) { - (IpAddr::V4(local_address), IpAddr::V4(remote_address)) => { - let (source, destination) = if info.direction == Direction::Send { - (local_address, remote_address) - } else { - (remote_address, local_address) - }; + match info.ip_addr() { + (IpAddr::V4(_), IpAddr::V4(_)) => { + let (source, destination) = info.ipvt_by_direction(); let header_size = HEADER_SIZE_IP4 + (32 / 8); @@ -185,11 +177,8 @@ impl Pcap { (ip.packet_size(), pnet_packet::ethernet::EtherTypes::Ipv4) } - (IpAddr::V6(local_address), IpAddr::V6(remote_address)) => { - let (source, destination) = match info.direction { - Direction::Send => (local_address, remote_address), - Direction::Receive => (remote_address, local_address), - }; + (IpAddr::V6(_), IpAddr::V6(_)) => { + let (source, destination) = info.ipvt_by_direction(); let mut ip = MutableIpv6Packet::new(buf).unwrap(); ip.set_version(6); diff --git a/crates/lib/src/capture/writer.rs b/crates/lib/src/capture/writer.rs index 476853c..1951e7e 100644 --- a/crates/lib/src/capture/writer.rs +++ b/crates/lib/src/capture/writer.rs @@ -1,22 +1,20 @@ -use std::io::Write; +use std::{io::Write, sync::Mutex}; -use crate::{ - capture::packet::{CapturePacket, Protocol}, - GDResult, +use super::{ + packet::{CapturePacket, Protocol}, + pcap::Pcap, }; - -use super::pcap::Pcap; - +use crate::GDResult; use lazy_static::lazy_static; -use std::sync::Mutex; lazy_static! { pub(crate) static ref CAPTURE_WRITER: Mutex>> = Mutex::new(None); } pub(crate) trait Writer { - fn write(&mut self, packet: &CapturePacket, data: &[u8]) -> crate::GDResult<()>; - fn new_connect(&mut self, packet: &CapturePacket) -> crate::GDResult<()>; + fn write(&mut self, packet: &CapturePacket, data: &[u8]) -> GDResult<()>; + fn new_connect(&mut self, packet: &CapturePacket) -> GDResult<()>; + //TODO: Close connection } impl Writer for Pcap { diff --git a/crates/lib/src/socket.rs b/crates/lib/src/socket.rs index f380c89..0ee26f6 100644 --- a/crates/lib/src/socket.rs +++ b/crates/lib/src/socket.rs @@ -4,10 +4,9 @@ use crate::{ GDResult, }; -use std::net::SocketAddr; use std::{ io::{Read, Write}, - net, + net::{self, SocketAddr}, }; const DEFAULT_PACKET_SIZE: usize = 1024; @@ -227,7 +226,7 @@ pub mod capture { } /// Represents the UDP protocol provider. - pub struct ProtocolUDP; + pub(crate) struct ProtocolUDP; impl ProtocolProvider for ProtocolUDP { fn protocol() -> Protocol { Protocol::UDP @@ -240,7 +239,7 @@ pub mod capture { /// * `I` - The inner socket type. /// * `P` - The protocol provider. #[derive(Clone, Debug)] - pub struct WrappedCaptureSocket { + pub(crate) struct WrappedCaptureSocket { inner: I, remote_address: SocketAddr, _protocol: PhantomData

, @@ -292,7 +291,7 @@ pub mod capture { /// /// # Returns /// A result indicating success or error in sending data. - fn send(&mut self, data: &[u8]) -> crate::GDResult<()> { + fn send(&mut self, data: &[u8]) -> GDResult<()> { let info = CapturePacket { direction: Direction::Send, protocol: P::protocol(),