tribufu/rust-gamedig
1
0
Fork 0
mirror of https://github.com/tribufu/rust-gamedig synced 2026-05-06 07:17:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

refactor: clean up

Browse source
This commit is contained in:
Cain 2024-01-13 09:18:56 +00:00 • committed by Douile
parent 3d47180e85
commit 0543cabce2
No known key found for this signature in database
GPG key ID: E048586A5FF6585C
6 changed files with 137 additions and 171 deletions
Download patch file Download diff file Expand all files Collapse all files

140
crates/cli/src/main.rs
View file

@ -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<u16>,
/// 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<u16>,
/// 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<String>,
/// 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<TimeoutSettings>,
/// Optional file path for packet capture file writer
#[cfg(feature = "packet_capture")]
#[arg(short, long)]
capture: Option<PathBuf>,
/// Optional extra settings for the server query.
#[command(flatten, next_help_heading = "Query options")]
extra_options: Option<ExtraRequestSettings>,
/// Optional timeout settings for the server query
#[command(flatten, next_help_heading = "Timeouts")]
timeout_settings: Option<TimeoutSettings>,
/// Optional extra settings for the server query
#[command(flatten, next_help_heading = "Query options")]
extra_options: Option<ExtraRequestSettings>,
},
/// 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<ExtraRequestSe
/// # Arguments
/// * `args` - A reference to the command line options.
/// * `result` - A reference to the result of the query.
fn output_result(args: &Cli, result: &dyn CommonResponse) {
match args.output_mode {
fn output_result(output: OutputMode, json: bool, result: &dyn CommonResponse) {
match output {
#[cfg(feature = "json")]
OutputMode::Generic if args.json => 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<R: serde::Serialize>(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(())
}

24
crates/lib/src/capture/mod.rs
View file

@ -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<String>) {
if let Some(file_name) = file_name {
pub fn setup_capture(file_path: Option<PathBuf>) {
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<String>) {
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<String>) {
/// 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<dyn Writer + Send + Sync>) {
crate::socket::capture::set_writer(writer);
}

94
crates/lib/src/capture/packet.rs
View file

@ -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<Self>;
}
@ -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<T: IpAddress>(&self) -> (T, T) {
pub(super) fn ipvt_by_direction<T: IpAddress>(&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<T>(&self, source: T, remote: T) -> (T, T) {
pub(self) const fn order<T>(&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::<Ipv6Addr>());
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();
}
}

23
crates/lib/src/capture/pcap.rs
View file

@ -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<W: Write> Pcap<W> {
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<W: Write> Pcap<W> {
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<W: Write> Pcap<W> {
(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);

18
crates/lib/src/capture/writer.rs
View file

@ -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<Option<Box<dyn Writer + Send + Sync>>> = 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<W: Write> Writer for Pcap<W> {

9
crates/lib/src/socket.rs
View file

@ -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<I, P> {
pub(crate) struct WrappedCaptureSocket<I, P> {
inner: I,
remote_address: SocketAddr,
_protocol: PhantomData<P>,
@ -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(),