641 lines
22 KiB
Rust
641 lines
22 KiB
Rust
use anyhow::{bail, ensure, Result};
|
|
use binrw::BinReaderExt;
|
|
use binrw::{BinRead, NullString};
|
|
use chacha20::cipher::KeyInit;
|
|
use chacha20::cipher::{KeyIvInit, StreamCipher, StreamCipherSeek};
|
|
use chacha20::ChaCha20;
|
|
use clap::Parser;
|
|
use dialoguer::theme::ColorfulTheme;
|
|
use dialoguer::Select;
|
|
use futures_util::FutureExt;
|
|
use poly1305::Poly1305;
|
|
use rand::{thread_rng, Rng};
|
|
use rhexdump::hexdump;
|
|
use rustyline_async::{Readline, ReadlineError, SharedWriter};
|
|
use std::collections::BTreeMap;
|
|
use std::error::Error;
|
|
use std::fmt::Display;
|
|
use std::io::Cursor;
|
|
use std::io::Write;
|
|
use std::iter;
|
|
use std::net::SocketAddr;
|
|
use std::net::ToSocketAddrs;
|
|
use std::net::{IpAddr, Ipv4Addr};
|
|
use std::path::PathBuf;
|
|
use std::time::{Duration, Instant};
|
|
use tokio::io::AsyncBufReadExt;
|
|
use tokio::net::UdpSocket;
|
|
use tokio::time;
|
|
|
|
mod hex_ii;
|
|
mod parser;
|
|
|
|
const KEY: &[u8; 32] = b"\x02\x04\x06\x08\x0a\x0c\x0e\x10\x12\x14\x16\x18\x1a\x1c\x1e\x20\x22\x24\x26\x28\x2a\x2c\x2e\x30\x32\x34\x36\x38\x3a\x3c\x3e\x40";
|
|
const INFO_PACKET: &[u8] = b"\x7f\x01\x00\x00\x07";
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct ServerFlags {
|
|
dedicated: bool,
|
|
force_vehicle: bool,
|
|
_rest: u8,
|
|
}
|
|
|
|
impl Display for ServerFlags {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
let force_vehicle = if self.force_vehicle { "F" } else { " " };
|
|
let dedicated = if self.dedicated { "D" } else { " " };
|
|
write!(f, "{}{}", force_vehicle, dedicated)?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl From<u8> for ServerFlags {
|
|
fn from(v: u8) -> Self {
|
|
ServerFlags {
|
|
dedicated: v & 0b1 != 0,
|
|
force_vehicle: v & 0b10 != 0,
|
|
_rest: (v & 0b11111100) >> 2,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, BinRead)]
|
|
#[br(little, magic = b"\xba\xce", import(rtt: Duration, addr: SocketAddr))]
|
|
pub struct Server {
|
|
#[br(calc=addr)]
|
|
addr: SocketAddr,
|
|
#[br(calc=rtt)]
|
|
rtt: Duration,
|
|
#[br(map = |v: (u8,u8)| format!("{}.{}",v.0,v.1))]
|
|
version: String,
|
|
port: u16,
|
|
max_players: u16,
|
|
cur_players: u16,
|
|
#[br(map = u8::into)]
|
|
flags: ServerFlags,
|
|
#[br(pad_size_to(0x20), map = |s :NullString| s.to_string())]
|
|
name: String,
|
|
#[br(pad_size_to(0x10), map = |s :NullString| s.to_string())]
|
|
mode: String,
|
|
#[br(pad_size_to(0x20), map = |s :NullString| s.to_string())]
|
|
map: String,
|
|
_pad: u8,
|
|
}
|
|
|
|
fn pad_copy(d: &[u8], l: usize) -> Vec<u8> {
|
|
let diff = d.len() % l;
|
|
if diff != 0 {
|
|
d.iter()
|
|
.copied()
|
|
.chain(iter::repeat(0).take(l - diff))
|
|
.collect()
|
|
} else {
|
|
d.to_vec()
|
|
}
|
|
}
|
|
|
|
fn pad(d: &mut Vec<u8>, l: usize) {
|
|
let diff = d.len() % l;
|
|
if diff != 0 {
|
|
d.extend(iter::repeat(0).take(l - diff))
|
|
}
|
|
}
|
|
|
|
struct Packet {
|
|
nonce: Vec<u8>,
|
|
data: Vec<u8>,
|
|
}
|
|
|
|
impl Packet {
|
|
fn encrypt(data: &[u8]) -> Packet {
|
|
let mut data: Vec<u8> = data.to_vec();
|
|
let mut rng = thread_rng();
|
|
let mut nonce = vec![0u8; 12];
|
|
rng.fill(nonce.as_mut_slice());
|
|
let mut cipher = ChaCha20::new(KEY.into(), nonce.as_slice().into());
|
|
cipher.seek(KEY.len() + 32);
|
|
cipher.apply_keystream(&mut data);
|
|
Packet { nonce, data }
|
|
}
|
|
|
|
fn get_tag(&self) -> Vec<u8> {
|
|
let mut sign_data = vec![];
|
|
sign_data.extend(pad_copy(&self.nonce, 16).iter());
|
|
sign_data.extend(pad_copy(&self.data, 16).iter());
|
|
sign_data.extend((self.nonce.len() as u64).to_le_bytes().iter());
|
|
sign_data.extend((self.data.len() as u64).to_le_bytes().iter());
|
|
let mut cipher = ChaCha20::new(KEY.into(), self.nonce.as_slice().into());
|
|
let mut poly_key = *KEY;
|
|
cipher.apply_keystream(&mut poly_key);
|
|
let signer = Poly1305::new(&poly_key.into());
|
|
signer.compute_unpadded(&sign_data).into_iter().collect()
|
|
}
|
|
|
|
fn bytes(&self) -> Vec<u8> {
|
|
let mut data = vec![];
|
|
data.extend(pad_copy(&self.nonce, 16).iter());
|
|
data.extend(pad_copy(&self.data, 16).iter());
|
|
data.extend((self.nonce.len() as u64).to_le_bytes().iter());
|
|
data.extend((self.data.len() as u64).to_le_bytes().iter());
|
|
data.extend(self.get_tag().iter());
|
|
data
|
|
}
|
|
|
|
fn decrypt(&self) -> Result<Vec<u8>> {
|
|
let mut data = self.data.clone();
|
|
let mut sign_data = data.clone();
|
|
pad(&mut sign_data, 16);
|
|
let mut nonce = self.nonce.clone();
|
|
pad(&mut nonce, 16);
|
|
let sign_data = nonce
|
|
.iter()
|
|
.chain(sign_data.iter())
|
|
.chain((self.nonce.len() as u64).to_le_bytes().iter())
|
|
.chain((self.data.len() as u64).to_le_bytes().iter())
|
|
.copied()
|
|
.collect::<Vec<u8>>();
|
|
let mut poly_key = *KEY;
|
|
let mut cipher = ChaCha20::new(KEY.into(), self.nonce.as_slice().into());
|
|
cipher.apply_keystream(&mut poly_key);
|
|
let signer = Poly1305::new(&poly_key.into());
|
|
let signature: Vec<u8> = signer.compute_unpadded(&sign_data).into_iter().collect();
|
|
|
|
if signature != self.get_tag() {
|
|
bail!("Invalid signature!");
|
|
};
|
|
cipher.seek(poly_key.len() + 32);
|
|
cipher.apply_keystream(&mut data);
|
|
Ok(data)
|
|
}
|
|
}
|
|
|
|
impl TryFrom<&[u8]> for Packet {
|
|
type Error = anyhow::Error;
|
|
fn try_from(data: &[u8]) -> Result<Self> {
|
|
let (mut nonce, data) = data.split_at(16);
|
|
let (mut data, tag) = data.split_at(data.len() - 16);
|
|
let nonce_len = u64::from_le_bytes(data[data.len() - 16..][..8].try_into()?) as usize;
|
|
let data_len = u64::from_le_bytes(data[data.len() - 8..].try_into()?) as usize;
|
|
data = &data[..data_len];
|
|
nonce = &nonce[..nonce_len];
|
|
let pkt = Packet {
|
|
nonce: nonce.into(),
|
|
data: data.into(),
|
|
};
|
|
if pkt.get_tag() != tag {
|
|
bail!("Invalid signature!");
|
|
}
|
|
Ok(pkt)
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub enum ServerEntry {
|
|
Alive(Server),
|
|
Dead { addr: SocketAddr, reason: String },
|
|
}
|
|
|
|
impl Display for ServerEntry {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
match self {
|
|
ServerEntry::Alive(srv) => write!(
|
|
f,
|
|
"[{}] {} ({} {}/{} Players on {}) version {} [{}] RTT: {:?}",
|
|
srv.addr,
|
|
srv.name,
|
|
srv.mode,
|
|
srv.cur_players,
|
|
srv.max_players,
|
|
srv.map,
|
|
srv.version,
|
|
srv.flags,
|
|
srv.rtt
|
|
),
|
|
ServerEntry::Dead { addr, reason } => write!(f, "[{}] (error: {})", addr, reason),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn encrypt(data: &[u8]) -> Vec<u8> {
|
|
Packet::encrypt(data).bytes()
|
|
}
|
|
|
|
fn decrypt(data: &[u8]) -> Result<Vec<u8>> {
|
|
Packet::try_from(data)?.decrypt()
|
|
}
|
|
|
|
async fn recv_from_timeout(
|
|
sock: &UdpSocket,
|
|
buf: &mut [u8],
|
|
timeout: f64,
|
|
) -> Result<(usize, SocketAddr)> {
|
|
Ok(time::timeout(Duration::from_secs_f64(timeout), sock.recv_from(buf)).await??)
|
|
}
|
|
|
|
async fn query_server<'a>(addr: SocketAddr) -> Result<Server> {
|
|
let mut buf = [0; 32 * 1024];
|
|
let socket = UdpSocket::bind("0.0.0.0:0").await?;
|
|
socket.connect(addr).await?;
|
|
let msg = encrypt(INFO_PACKET);
|
|
let t_start = Instant::now();
|
|
socket.send(&msg).await?;
|
|
let size = recv_from_timeout(&socket, &mut buf, 5.0).await?.0;
|
|
let rtt = t_start.elapsed();
|
|
let data = decrypt(&buf[..size])?;
|
|
if !data.starts_with(&[0xba, 0xce]) {
|
|
// Server Info
|
|
bail!("Invalid response");
|
|
}
|
|
let mut cur = Cursor::new(&data);
|
|
let info: Server = cur.read_le_args((rtt, addr))?;
|
|
if info.port != addr.port() {
|
|
eprint!("[WARN] Port differs for {}: {}", addr, info.port);
|
|
}
|
|
if cur.position() != (data.len() as u64) {
|
|
bail!("Leftover data");
|
|
}
|
|
Ok(info)
|
|
}
|
|
|
|
async fn get_servers(master_addr: &str) -> Result<(Duration, Vec<ServerEntry>)> {
|
|
let master_addr: SocketAddr = master_addr.to_socket_addrs()?.next().unwrap();
|
|
let mut rtt = std::time::Duration::from_secs_f32(0.0);
|
|
let mut servers = vec![];
|
|
let mut buf = [0; 32 * 1024];
|
|
let master = UdpSocket::bind("0.0.0.0:0").await?;
|
|
master.connect(master_addr).await?;
|
|
for n in 0..(256 / 32) {
|
|
let data = format!("Brw={},{}\0", n * 32, (n + 1) * 32);
|
|
let data = &encrypt(data.as_bytes());
|
|
let t_start = Instant::now();
|
|
master.send(data).await?;
|
|
let size = master.recv(&mut buf).await?;
|
|
rtt += t_start.elapsed();
|
|
let data = decrypt(&buf[..size])?;
|
|
if data.starts_with(b"\0\0\0\0}") {
|
|
for chunk in data[5..].chunks(6) {
|
|
if chunk.iter().all(|v| *v == 0) {
|
|
break;
|
|
}
|
|
let port = u16::from_le_bytes(chunk[chunk.len() - 2..].try_into()?);
|
|
let addr = SocketAddr::from(([chunk[0], chunk[1], chunk[2], chunk[3]], port));
|
|
let server = match query_server(addr).await {
|
|
Ok(server) => ServerEntry::Alive(server),
|
|
Err(err) => ServerEntry::Dead {
|
|
addr,
|
|
reason: err.to_string(),
|
|
},
|
|
};
|
|
servers.push(server);
|
|
}
|
|
}
|
|
}
|
|
rtt = Duration::from_secs_f64(rtt.as_secs_f64() / ((256 / 32) as f64));
|
|
Ok((rtt, servers))
|
|
}
|
|
|
|
fn indent_hexdump(data: &[u8], indentation: usize, label: &str) -> String {
|
|
let mut out = String::new();
|
|
let indent = " ".repeat(indentation);
|
|
out.push_str(&indent);
|
|
out.push_str(label);
|
|
out.push('\n');
|
|
for line in rhexdump::hexdump(data).lines() {
|
|
out.push_str(&indent);
|
|
out.push_str(line);
|
|
out.push('\n');
|
|
}
|
|
out.trim_end().to_owned()
|
|
}
|
|
|
|
#[derive(Default, Debug)]
|
|
struct State {
|
|
client: BTreeMap<usize, BTreeMap<u8, usize>>,
|
|
server: BTreeMap<usize, BTreeMap<u8, usize>>,
|
|
}
|
|
|
|
impl State {
|
|
fn update_client(&mut self, data: &[u8]) {
|
|
data.iter().enumerate().for_each(|(pos, b)| {
|
|
*self.client.entry(pos).or_default().entry(*b).or_default() += 1;
|
|
});
|
|
}
|
|
fn update_server(&mut self, data: &[u8]) {
|
|
data.iter().enumerate().for_each(|(pos, b)| {
|
|
*self.server.entry(pos).or_default().entry(*b).or_default() += 1;
|
|
});
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, PartialEq, Eq, Copy, Clone)]
|
|
enum Direction {
|
|
Client,
|
|
Server,
|
|
Both,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
enum CmdResult {
|
|
Exit,
|
|
Packet {
|
|
data: Vec<u8>,
|
|
direction: Direction,
|
|
},
|
|
Fuzz {
|
|
direction: Direction,
|
|
start: usize,
|
|
end: usize,
|
|
chance: (u32, u32),
|
|
},
|
|
NoFuzz,
|
|
Log(bool),
|
|
}
|
|
|
|
async fn handle_line(
|
|
line: &str,
|
|
state: &State,
|
|
stdout: &mut SharedWriter,
|
|
) -> Result<Option<CmdResult>> {
|
|
use CmdResult::*;
|
|
let cmd: Vec<&str> = line.trim().split_ascii_whitespace().collect();
|
|
match cmd[..] {
|
|
["log", "off"] => Ok(Some(Log(false))),
|
|
["log", "on"] => Ok(Some(Log(true))),
|
|
["state", pos] => {
|
|
let pos = pos.parse()?;
|
|
writeln!(stdout, "Client: {:?}", state.client.get(&pos))?;
|
|
writeln!(stdout, "Server: {:?}", state.server.get(&pos))?;
|
|
Ok(None)
|
|
}
|
|
[dir @ ("client" | "server"), ref args @ ..] => {
|
|
let mut data: Vec<u8> = vec![];
|
|
for args in args.iter() {
|
|
let args = hex::decode(args)?;
|
|
data.extend(args);
|
|
}
|
|
Ok(Some(CmdResult::Packet {
|
|
data,
|
|
direction: match dir {
|
|
"client" => Direction::Client,
|
|
"server" => Direction::Server,
|
|
_ => unreachable!(),
|
|
},
|
|
}))
|
|
}
|
|
["fuzz", dir @ ("client" | "server" | "both"), start, end, chance_num, chance_den] => {
|
|
let direction = match dir {
|
|
"client" => Direction::Client,
|
|
"server" => Direction::Server,
|
|
"both" => Direction::Both,
|
|
_ => unreachable!(),
|
|
};
|
|
let start = start.parse()?;
|
|
let end = end.parse()?;
|
|
if start > end {
|
|
bail!("Fuzz start>end");
|
|
}
|
|
let res = CmdResult::Fuzz {
|
|
direction,
|
|
start,
|
|
end,
|
|
chance: (chance_num.parse()?, chance_den.parse()?),
|
|
};
|
|
Ok(Some(res))
|
|
}
|
|
["fuzz", "off"] => Ok(Some(CmdResult::NoFuzz)),
|
|
["exit"] => Ok(Some(CmdResult::Exit)),
|
|
[""] => Ok(None),
|
|
_ => bail!("Unknown command: {:?}", line),
|
|
}
|
|
}
|
|
|
|
async fn run_proxy(
|
|
remote_addr: &SocketAddr,
|
|
local_addr: &SocketAddr,
|
|
logfile: &Option<PathBuf>,
|
|
) -> Result<()> {
|
|
let mut print_log = false;
|
|
let mut state = State::default();
|
|
let mut logfile = match logfile {
|
|
Some(path) => Some(std::fs::File::create(path)?),
|
|
None => None,
|
|
};
|
|
let mut fuzz = None;
|
|
let mut rng = thread_rng();
|
|
let mut client_addr: Option<SocketAddr> = None;
|
|
let local = UdpSocket::bind(local_addr).await?;
|
|
let remote = UdpSocket::bind("0.0.0.0:0").await?;
|
|
remote.connect(remote_addr).await?;
|
|
let mut local_buf = vec![0; 32 * 1024];
|
|
let mut remote_buf = vec![0; 32 * 1024];
|
|
println!("Proxy listening on {}", local_addr);
|
|
let (mut rl, mut stdout) = Readline::new(format!("{}> ", remote_addr)).unwrap();
|
|
loop {
|
|
tokio::select! {
|
|
line = rl.readline().fuse() => {
|
|
match line {
|
|
Ok(line) => {
|
|
let line=line.trim();
|
|
rl.add_history_entry(line.to_owned());
|
|
match handle_line(line, &state, &mut stdout).await {
|
|
Ok(Some(result)) => {
|
|
match result {
|
|
CmdResult::Packet{data,direction} => {
|
|
let data=encrypt(&data);
|
|
match direction {
|
|
Direction::Client => {
|
|
if client_addr.is_some() {
|
|
local
|
|
.send_to(&data, client_addr.unwrap())
|
|
.await?;
|
|
} else {
|
|
writeln!(stdout,"Error: No client address")?;
|
|
}
|
|
},
|
|
Direction::Server => {
|
|
remote.send(&data).await?;
|
|
}
|
|
Direction::Both => unreachable!()
|
|
};
|
|
}
|
|
CmdResult::Log(log) => {
|
|
print_log=log;
|
|
}
|
|
CmdResult::Exit => break Ok(()),
|
|
CmdResult::NoFuzz => {
|
|
fuzz=None;
|
|
}
|
|
CmdResult::Fuzz { .. } => {
|
|
fuzz=Some(result)
|
|
},
|
|
}
|
|
},
|
|
Ok(None) => (),
|
|
Err(msg) => {
|
|
writeln!(stdout, "Error: {}", msg)?;
|
|
}
|
|
}
|
|
},
|
|
Err(ReadlineError::Eof) =>{ writeln!(stdout, "Exiting...")?; break Ok(()) },
|
|
Err(ReadlineError::Interrupted) => {
|
|
writeln!(stdout, "^C")?;
|
|
break Ok(());
|
|
},
|
|
Err(err) => {
|
|
writeln!(stdout, "Received err: {:?}", err)?;
|
|
writeln!(stdout, "Exiting...")?;
|
|
break Ok(());
|
|
}
|
|
}
|
|
}
|
|
local_res = local.recv_from(&mut local_buf) => {
|
|
let (size, addr) = local_res?;
|
|
client_addr.get_or_insert(addr);
|
|
let mut data = Packet::try_from(&local_buf[..size])?.decrypt()?;
|
|
state.update_client(&data);
|
|
if print_log {
|
|
writeln!(stdout,"{}", indent_hexdump(&data, 0, &format!("OUT: {}", addr)))?;
|
|
}
|
|
if let Some(lf) = logfile.as_mut() {
|
|
writeln!(lf, ">{:?} {} {}", addr, data.len(), hex::encode(&data))?;
|
|
};
|
|
if let Some(CmdResult::Fuzz{direction,start,end,chance}) = fuzz {
|
|
if (direction==Direction::Server || direction==Direction::Both) && rng.gen_ratio(chance.0,chance.1) {
|
|
rng.fill(&mut data[start..end]);
|
|
}
|
|
}
|
|
remote.send(&encrypt(&data)).await?;
|
|
}
|
|
remote_res = remote.recv_from(&mut remote_buf) => {
|
|
let (size, addr) = remote_res?;
|
|
let mut data = Packet::try_from(&remote_buf[..size])?.decrypt()?;
|
|
state.update_server(&data);
|
|
if print_log {
|
|
writeln!(stdout,"\r{}", indent_hexdump(&data, 5, &format!("IN: {}", addr)))?;
|
|
}
|
|
if let Some(lf) = logfile.as_mut() {
|
|
writeln!(lf, "<{:?} {} {}", addr, data.len(), hex::encode(&data))?;
|
|
};
|
|
if client_addr.is_some() {
|
|
if let Some(CmdResult::Fuzz{direction,start,end,chance}) = &fuzz {
|
|
if (*direction==Direction::Client || *direction==Direction::Both) && rng.gen_ratio(chance.0,chance.1) {
|
|
rng.fill(&mut data[*start..*end]);
|
|
}
|
|
}
|
|
local
|
|
.send_to(&encrypt(&data), client_addr.unwrap())
|
|
.await?;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn send_master_cmd(sock: &UdpSocket, cmd: &str) -> Result<Vec<u8>> {
|
|
let mut buf = [0; 32 * 1024];
|
|
let mut data: Vec<u8> = cmd.as_bytes().to_vec();
|
|
data.push(0);
|
|
let data = &encrypt(&data);
|
|
sock.send(data).await?;
|
|
let size = recv_from_timeout(sock, &mut buf, 5.0).await?.0;
|
|
decrypt(&buf[..size])
|
|
}
|
|
|
|
async fn run_master_shell(master_addr: &str) -> Result<()> {
|
|
let master = UdpSocket::bind("0.0.0.0:0").await?;
|
|
master.connect(master_addr).await?;
|
|
let (mut rl, mut stdout) = Readline::new(format!("{}> ", master_addr)).unwrap();
|
|
loop {
|
|
tokio::select! {
|
|
line = rl.readline().fuse() => {
|
|
match line {
|
|
Ok(line) => {
|
|
let line=line.trim();
|
|
rl.add_history_entry(line.to_owned());
|
|
writeln!(stdout,"[CMD] {line}")?;
|
|
match send_master_cmd(&master,line).await {
|
|
Ok(data) => writeln!(stdout,"{}",hexdump(&data))?,
|
|
Err(e) => writeln!(stdout,"Error: {e}")?
|
|
}
|
|
}
|
|
Err(ReadlineError::Eof) =>{ writeln!(stdout, "Exiting...")?; break Ok(()) },
|
|
Err(ReadlineError::Interrupted) => {
|
|
writeln!(stdout, "^C")?;
|
|
break Ok(());
|
|
},
|
|
Err(err) => {
|
|
writeln!(stdout, "Receive error: {err}")?;
|
|
break Err(err.into());
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Parser, Debug)]
|
|
#[clap(author, version, about, long_about = None)]
|
|
struct Args {
|
|
/// Server to connect to (if unspecified will query the master server)
|
|
server: Option<SocketAddr>,
|
|
/// Only list servers without starting proxy
|
|
#[clap(short, long, action)]
|
|
list: bool,
|
|
/// Local Address to bind to
|
|
#[clap(short,long, default_value_t = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 28086))]
|
|
addr: SocketAddr,
|
|
/// Master server to query for running games
|
|
#[clap(short, long, default_value = "scrapland.mercurysteam.com:5000")]
|
|
master: String,
|
|
/// Path of file to log decrypted packets to
|
|
#[clap(short = 'f', long)]
|
|
logfile: Option<PathBuf>,
|
|
}
|
|
|
|
#[tokio::main]
|
|
async fn main() -> Result<()> {
|
|
let args = Args::parse();
|
|
if args.list && args.server.is_some() {
|
|
let addr = args.server.unwrap();
|
|
let server = match query_server(addr).await {
|
|
Ok(server) => ServerEntry::Alive(server),
|
|
Err(msg) => ServerEntry::Dead {
|
|
addr,
|
|
reason: msg.to_string(),
|
|
},
|
|
};
|
|
println!("{}", server);
|
|
return Ok(());
|
|
}
|
|
if let Some(server) = args.server {
|
|
run_proxy(&server, &args.addr, &args.logfile).await?;
|
|
return Ok(());
|
|
}
|
|
loop {
|
|
let (rtt, servers) = get_servers(&args.master).await?;
|
|
println!("Master RTT: {:?}", rtt);
|
|
if args.list {
|
|
for server in servers {
|
|
println!("{}", server);
|
|
}
|
|
return Ok(());
|
|
}
|
|
let selection = Select::with_theme(&ColorfulTheme::default())
|
|
.items(&servers)
|
|
.with_prompt("Select server (press Esc to drop into master server command shell)")
|
|
.interact_opt()?
|
|
.map(|v| &servers[v]);
|
|
match selection {
|
|
Some(ServerEntry::Dead { addr, reason }) => {
|
|
eprintln!("{:?} returned an error: {}", addr, reason)
|
|
}
|
|
Some(ServerEntry::Alive(srv)) => {
|
|
return run_proxy(&srv.addr, &args.addr, &args.logfile).await;
|
|
}
|
|
None => {
|
|
return run_master_shell(&args.master).await;
|
|
}
|
|
}
|
|
}
|
|
}
|