1 files changed, 314 insertions, 0 deletions
diff --git a/src/bin/mcquery.rs b/src/bin/mcquery.rs
new file mode 100644
index 0000000..8dcd1bc
--- /dev/null
+++ b/src/bin/mcquery.rs
@@ -0,0 +1,314 @@
+extern crate getopts;
+use getopts::Options;
+use std::env;
+use rand::Rng;
+use std::net::{UdpSocket,Ipv4Addr,ToSocketAddrs,SocketAddr};
+
+pub mod querynet {
+    use std::net::{UdpSocket, SocketAddr};
+    use std::str;
+    use bytes::{BytesMut, BufMut};
+    use byteorder::{ByteOrder,LittleEndian};
+    use std::collections::HashMap;
+    use json::object;
+
+    pub struct BasicStat {
+        pub motd: String,
+        pub gametype: String,
+        pub map: String,
+        pub numplayers: u32,
+        pub maxplayers: u32,
+        pub hostport: u16,
+        pub hostip: String
+    }
+
+    impl BasicStat {
+        pub fn ip(&self) {
+            println!("{}:{}", self.hostip, self.hostport);
+        }
+
+        pub fn players(&self) {
+            println!("{}/{}", self.numplayers, self.maxplayers);
+        }
+
+        pub fn info(&self) {
+            println!("{}\n{}, {}", self.motd, self.map, self.gametype)
+        }
+
+        pub fn json(&self) -> String {
+            object!{
+                motd: self.motd.clone(),
+                gametype: self.gametype.clone(),
+                map: self.map.clone(),
+                players: object!{
+                    max: self.maxplayers,
+                    online: self.numplayers,
+                },
+                hostport: self.hostport,
+                hostip: self.hostip.clone()
+            }.dump()
+        }
+
+    }
+
+    pub struct FullStat {
+        pub motd: String,
+        pub gametype: String,
+        pub game_id: String,
+        pub version: String,
+        pub plugins: String, // later?
+        pub map: String,
+        pub numplayers: u32,
+        pub maxplayers: u32,
+        pub hostport: u16,
+        pub hostip: String,
+        pub players: Vec<String>,
+    }
+
+    impl FullStat {
+
+        pub fn set_players(&mut self, players: Vec<String>) {
+            self.players = players;
+        }
+
+        pub fn json(&self) -> String {
+            object!{
+                motd: self.motd.clone(),
+                gametype: self.gametype.clone(),
+                gameid: self.game_id.clone(),
+                version: self.version.clone(),
+                plugins: self.plugins.clone(),
+                map: self.map.clone(),
+                players: object!{
+                    max: self.maxplayers,
+                    online: self.numplayers,
+                    list: self.players.clone()
+                },
+                hostport: self.hostport,
+                hostip: self.hostip.clone()
+            }.dump()
+        }
+    }
+
+    pub fn convert_kv(kv: HashMap<String, String>) -> FullStat {
+        FullStat {
+            motd: match kv.get(&String::from("hostname")) { Some(v) => v.clone(), _ => String::from("") },
+            gametype: match kv.get(&String::from("gametype")) { Some(v) => v.clone(), _ => String::from("") },
+            game_id: match kv.get(&String::from("game_id")) { Some(v) => v.clone(), _ => String::from("") },
+            version: match kv.get(&String::from("version")) { Some(v) => v.clone(), _ => String::from("") },
+            plugins: match kv.get(&String::from("plugins")) { Some(v) => v.clone(), _ => String::from("") },
+            map: match kv.get(&String::from("map")) { Some(v) => v.clone(), _ => String::from("") },
+            numplayers: match kv.get(&String::from("numplayers")) { Some(v) => v.clone().parse().expect("NaN"), _ => 0 },
+            maxplayers: match kv.get(&String::from("maxplayers")) { Some(v) => v.clone().parse().expect("NaN"), _ => 0 },
+            hostport: match kv.get(&String::from("hostport")) { Some(v) => v.clone().parse().expect("NaN"), _ => 0 },
+            hostip: match kv.get(&String::from("hostip")) { Some(v) => v.clone(), _ => String::from("") },
+            players: Vec::new()
+        }
+    }
+
+
+    pub fn slice_to_string(slice: &[u8]) -> String {
+        String::from_utf16(&slice.iter().map(|&slice| slice as u16).collect::<Vec<u16>>()[..]).expect("String is not UTF-16, for some reason.")
+    }
+
+    pub fn send_packet(socket: &UdpSocket, addr: SocketAddr, ptype: u8, session_id: u32, payload: &[u8]) {
+        let mut packet = BytesMut::with_capacity(16384);
+        packet.put(&b"\xFE\xFD"[..]);
+        packet.put_u8(ptype);
+        packet.put_u32(session_id);
+        packet.put(payload);
+        socket.send_to(&packet, addr).expect("Couldn't send packet.");
+    }
+
+    pub fn recv_packet(socket: &UdpSocket) -> BytesMut {
+        let mut raw_packet = [0; 16384];
+        let (packet_length, _src_addr) = socket.recv_from(&mut raw_packet).expect("Drring, drring.");
+        let mut packet = BytesMut::with_capacity(16384);
+        packet.put_slice(&raw_packet[5..packet_length]);
+        return packet;
+    }
+
+    pub fn send_handshake(socket: &UdpSocket, session_id: u32, addr: SocketAddr) {
+        send_packet(socket, addr, 9, session_id, &[0 as u8; 0]);
+    }
+
+    pub fn recv_handshake(socket: &UdpSocket) -> u32 {
+        str::from_utf8(&recv_packet(socket)).expect("Non-UTF8 string.").trim_end_matches(char::from(0)).parse::<u32>().expect("Non-number string.")
+    }
+
+    pub fn send_basicstat(socket: &UdpSocket, session_id: u32, addr: SocketAddr, challenge: u32) {
+        let mut challenge_token = BytesMut::with_capacity(32);
+        challenge_token.put_u32(challenge);
+        send_packet(socket, addr, 0, session_id, &challenge_token);
+    }
+
+    pub fn recv_basicstat(socket: &UdpSocket) -> BasicStat {
+        let bs_buffer = &recv_packet(socket)[..];
+        let bs_vector: Vec<&[u8]> = bs_buffer.split(|&ch| ch == 0).collect::<Vec<&[u8]>>();
+        BasicStat {
+            motd: slice_to_string(bs_vector[0]),
+            gametype: slice_to_string(bs_vector[1]),
+            map: slice_to_string(bs_vector[2]),
+            numplayers: str::from_utf8(bs_vector[3]).expect("Players are not UTF-8.").parse().expect("Not a number."),
+            maxplayers: str::from_utf8(bs_vector[4]).expect("Players are not UTF-8.").parse().expect("Not a number."),
+            hostport: LittleEndian::read_u16(&bs_vector[5][..2]),
+            hostip: slice_to_string(&bs_vector[5][2..]),
+        }
+    }
+
+    pub fn send_fullstat(socket: &UdpSocket, session_id: u32, addr: SocketAddr, challenge: u32) {
+        let mut challenge_token = BytesMut::with_capacity(64);
+        challenge_token.put_u32(challenge);
+        challenge_token.put_u32(0);
+        send_packet(socket, addr, 0, session_id, &challenge_token);
+    }
+
+    pub fn recv_fullstat(socket: &UdpSocket) -> FullStat {
+        let fs_buffer = &recv_packet(socket)[11..];
+        let mut fs_vector: Vec<&[u8]> = fs_buffer.split(|&ch| ch == 0).collect::<Vec<&[u8]>>();
+        let mut fullset_kv: HashMap<String, String> = HashMap::new();
+        let mut players: Vec<String> = Vec::new();
+        while fs_vector[0] != [0;0] {
+            let key = fs_vector.remove(0);
+            let value = fs_vector.remove(0);
+            fullset_kv.insert(slice_to_string(key), slice_to_string(value));
+        }
+        fs_vector.remove(0);
+        fs_vector.remove(0); // padding moment
+        fs_vector.remove(0);
+        while fs_vector[0] != [0;0] {
+            players.push(slice_to_string(fs_vector.remove(0)));
+        }
+        let mut fullstat = convert_kv(fullset_kv);
+        fullstat.set_players(players);
+        return fullstat;
+    }
+
+}
+
+fn usage(program: &str, opts: Options) {
+    let brief = format!("Usage: {} ADDRESS[:PORT]", program);
+    print!("{}", opts.usage(&brief));
+}
+
+fn main() {
+    let args: Vec<String> = env::args().collect();
+    let program = args[0].clone();
+    let mut opts = Options::new();
+    opts.optflag("h", "help", "display this help and exit");
+    opts.optflag("b", "basicstats", "sends a basic stat request");
+    let flags = match opts.parse(&args[1..]) {
+        Ok(m) => { m }
+        Err(f) => { panic!("{}", f.to_string()) }
+    };
+    if flags.opt_present("h") {
+        usage(&program, opts);
+        return;
+    }
+    let basic_stat = flags.opt_present("b");
+    let server = if !flags.free.is_empty() {
+        if flags.free[0].contains(":") {
+            flags.free[0].clone()
+        } else {
+            flags.free[0].clone() + ":25565"
+        }
+    } else {
+        usage(&program, opts);
+        return;
+    };
+
+    let mut rng = rand::thread_rng();
+    let session_id = rng.gen::<u32>() & 0x0F0F0F0F;
+
+    let addrs: Vec<SocketAddr> = server.to_socket_addrs().expect("Unable to resolve domain.").collect();
+    let socket = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 0)).expect("Couldn't bind to address.");
+    let addr = addrs[0];
+
+    querynet::send_handshake(&socket, session_id, addr);
+    let challenge_token = querynet::recv_handshake(&socket);
+
+    if basic_stat {
+        querynet::send_basicstat(&socket, session_id, addr, challenge_token);
+        let basic_stats = querynet::recv_fullstat(&socket);
+        println!("{}", basic_stats.json());
+    } else {
+        querynet::send_fullstat(&socket, session_id, addr, challenge_token);
+        let full_stats = querynet::recv_fullstat(&socket);
+        println!("{}", full_stats.json());
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    static REQUEST_HS: [u8; 7] = [0xFE, 0xFD, 0x09, 0x00, 0x00, 0x0A, 0xFE];
+    static RESPONSE_HS: [u8; 13] = [0x09, 0x00, 0x00, 0x0A, 0xFE, 0x39, 0x35, 0x31, 0x33, 0x33, 0x30, 0x37, 0x00];
+
+    // #[test]
+    // fn test_encode() {
+        // assert_eq!(&querynet::encode_packet(9, 2814, &[0 as u8;0])[..], REQUEST_HS);
+    // }
+
+    // #[test]
+    // fn test_decode() {
+        // let packet = &RESPONSE_HS[..];
+        // assert_eq!(querynet::decode_packet(packet), b"\x39\x35\x31\x33\x33\x30\x37\x00");
+    // }
+
+    #[test]
+    fn test_send_hs() {
+        let addr = SocketAddr::from(([127, 0, 0, 1], 5005));
+        let socket = UdpSocket::bind(addr).expect("Couldn't bind to address.");
+        querynet::send_handshake(&socket, 2814, addr);
+        let mut handshake: [u8; 7] = [0 as u8 ;7];
+        socket.recv_from(&mut handshake).expect("Didn't receive data.");
+        assert_eq!(handshake, REQUEST_HS);
+    }
+
+    #[test]
+    fn test_receive_hs() {
+        let addr = SocketAddr::from(([127, 0, 0, 1], 5006));
+        let socket = UdpSocket::bind(addr).expect("Couldn't bind to address.");
+        socket.send_to(&RESPONSE_HS[..], addr).expect("Didn't send");
+        assert_eq!(querynet::recv_handshake(&socket), 9513307);
+    }
+
+    #[test]
+    fn test_hs() {
+        let addr = SocketAddr::from(([51, 75, 186, 103], 25625)); // tarkoza lmao
+        let socket = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 0)).expect("Couldn't bind to address.");
+        let mut rng = rand::thread_rng();
+        let session_id = rng.gen::<u32>() & 0x0F0F0F0F;
+        println!("SESSION_ID: {} ", session_id);
+        querynet::send_handshake(&socket, session_id, addr);
+        let ch_tk = querynet::recv_handshake(&socket);
+        println!("CHALLENGE_TOKEN: {}", ch_tk);
+    }
+
+    #[test]
+    fn test_basicstat() {
+        let addr = SocketAddr::from(([51, 75, 186, 103], 25625)); // tarkoza lmao
+        let socket = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 0)).expect("Couldn't bind to address.");
+        let mut rng = rand::thread_rng();
+        let session_id = rng.gen::<u32>() & 0x0F0F0F0F;
+        querynet::send_handshake(&socket, session_id, addr);
+        let challenge_token = querynet::recv_handshake(&socket);
+        querynet::send_basicstat(&socket, session_id, addr, challenge_token);
+        let basic_stats = querynet::recv_basicstat(&socket);
+        assert_eq!(basic_stats.map, "Tkz");
+    }
+
+    #[test]
+    fn test_fullstat() {
+        let addr = SocketAddr::from(([51, 75, 186, 103], 25625)); // tarkoza lmao
+        let socket = UdpSocket::bind((Ipv4Addr::UNSPECIFIED, 0)).expect("Couldn't bind to address.");
+        let mut rng = rand::thread_rng();
+        let session_id = rng.gen::<u32>() & 0x0F0F0F0F;
+        querynet::send_handshake(&socket, session_id, addr);
+        let challenge_token = querynet::recv_handshake(&socket);
+        querynet::send_fullstat(&socket, session_id, addr, challenge_token);
+        let full_stats = querynet::recv_fullstat(&socket);
+        assert_eq!(full_stats.map, "Tkz");
+    }
+}