rcon.rs

// [dependencies]
// pico-args = "0.5.0"
// regex = "1.7.0"

use std::io;
use std::io::Read;
use std::io::Write;
use std::net::TcpStream;

use regex::Regex;

type Result<T> = std::io::Result<T>;

const HEADER_SIZE: usize = 10;

#[derive(Debug)]
struct Arguments {
    pub host: Option<String>,
    pub password: String,
    pub command: Option<String>,
}

#[derive(Debug)]
struct Packet {
    pub id: i32,
    pub kind: i32,
    pub payload: String,
}

impl Packet {
    pub fn new(id: i32, kind: i32) -> Self {
        Self {
            id,
            kind,
            payload: String::new(),
        }
    }
}

trait RconRead {
    fn read_packet(&mut self) -> Result<Packet>;

    fn read_i32(&mut self) -> Result<i32>;

    fn read_string(&mut self, len: usize) -> Result<String>;
}

impl<R> RconRead for R
where
    R: Read,
{
    fn read_packet(&mut self) -> Result<Packet> {
        let len = self.read_i32()?;
        let id = self.read_i32()?;
        let kind = self.read_i32()?;
        let payload = self.read_string((len as usize).saturating_sub(HEADER_SIZE))?;
        let mut end = [0u8; 2];
        self.read(&mut end)?;
        Ok(Packet { id, kind, payload })
    }

    fn read_i32(&mut self) -> Result<i32> {
        let mut buffer = [0u8; 4];
        self.read(&mut buffer)?;
        Ok(i32::from_le_bytes(buffer))
    }

    fn read_string(&mut self, len: usize) -> Result<String> {
        let mut buffer = vec![0u8; len];
        self.read_exact(&mut buffer)?;
        Ok(String::from_utf8(buffer).unwrap())
    }
}

trait RconWrite {
    fn write_packet(&mut self, packet: Packet) -> Result<()>;

    fn write_i32(&mut self, int: i32) -> Result<()>;

    fn write_string(&mut self, string: String) -> Result<()>;
}

impl<W> RconWrite for W
where
    W: Write,
{
    fn write_packet(&mut self, packet: Packet) -> Result<()> {
        let mut bytes = Vec::new();
        let len = (HEADER_SIZE + packet.payload.len()) as i32;
        bytes.write_i32(len)?;
        bytes.write_i32(packet.id)?;
        bytes.write_i32(packet.kind)?;
        bytes.write_string(packet.payload)?;
        bytes.write(&[0, 0])?;
        self.write_all(&bytes)?;
        Ok(())
    }

    fn write_i32(&mut self, int: i32) -> Result<()> {
        self.write(&int.to_le_bytes())?;
        Ok(())
    }

    fn write_string(&mut self, string: String) -> Result<()> {
        self.write_all(string.as_bytes())?;
        Ok(())
    }
}

fn main() -> Result<()> {
    let args = parse_args().unwrap();
    let mut id = 0;
    let mut stream = TcpStream::connect(args.host.unwrap_or("localhost:25575".to_string()))?;
    let mut packet = Packet::new(id, 3);
    packet.payload = args.password;
    stream.write_packet(packet)?;
    let packet = stream.read_packet()?;
    if packet.id != id {
        if packet.id == -1 {
            println!("Authentication failed due to invalid password!");
        } else {
            println!("Unexpected packet response! Packet id: {}", packet.id);
        }
        return Ok(());
    }
    if let Some(command) = args.command {
        id += 1;
        let mut packet = Packet::new(id, 2);
        packet.payload = command;
        stream.write_packet(packet)?;
        let packet = stream.read_packet()?;
        let message = escape_mc_format(packet.payload);
        if !message.is_empty() {
            println!("{}", message);
        }
    } else {
        println!("Use 'quit' or Ctrl-C to exit.");
        let stdin = io::stdin();
        loop {
            let mut command = String::new();
            print!("\x1b[33m>\x1b[0m ");
            io::stdout().flush()?;
            stdin.read_line(&mut command)?;
            command = command.trim().to_string();
            if command.eq("quit") {
                break;
            }
            id += 1;
            let mut packet = Packet::new(id, 2);
            packet.payload = command;
            stream.write_packet(packet)?;
            let packet = stream.read_packet()?;
            let message = escape_mc_format(packet.payload);
            if !message.is_empty() {
                println!("{}", message);
            }
        }
    }
    Ok(())
}

fn parse_args() -> std::result::Result<Arguments, pico_args::Error> {
    let mut args = pico_args::Arguments::from_env();
    Ok(Arguments {
        host: args.opt_value_from_str("-h")?,
        password: args.value_from_str("-p")?,
        command: args.opt_value_from_str("-c")?,
    })
}

fn escape_mc_format(string: String) -> String {
    let re = Regex::new(r"(§[a-z0-9]|\n$)").unwrap();
    re.replace_all(string.as_str(), "").to_string()
}