path: root/src/bin/mcquery.rs

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");
    }
}