use std::net::{ToSocketAddrs, UdpSocket};
use std::{io::{self, stdin, stdout, Read, Write}, net::{IpAddr, SocketAddr, TcpListener, TcpStream}};
use clap::{ArgAction, Parser, Subcommand};
#[derive(Parser)]
#[command(version, about)]
#[command(propagate_version = true)]
pub struct Config {
#[clap(long, short, action, action=ArgAction::SetFalse)]
tcp: bool,
#[clap(long, short, action)]
udp: bool,
#[command(subcommand)]
command: Op,
}
#[derive(Subcommand)]
enum Op {
Send { host: String, port: u16 },
Listen {
port: u16,
}
}
fn write_all(mut from: impl Read, mut to: impl Write) -> io::Result<usize> {
let mut buf = [0_u8; 1024];
let mut total = 0;
while let Ok(n) = from.read(&mut buf) {
if n == 0 { break; }
total += n;
to.write_all(&buf[0..n])?;
}
Ok(total)
}
fn tcp(conf: Config) -> crate::Result<()> {
match conf.command {
Op::Send { host, port } => {
let stream = TcpStream::connect(format!("{}:{}",host,port))?;
let stdin = stdin().lock();
write_all(stdin, stream)?;
}
Op::Listen { port } => {
let ip = IpAddr::V6([0_u8; 16].into());
let socket = SocketAddr::new(ip, port);
let stream = TcpListener::bind(socket)?;
let mut stdout = stdout().lock();
println!("Listening on {socket}");
let stream = stream.incoming().next().unwrap()?;
write_all(stream, &mut stdout)?;
}
}
Ok(())
}
fn udp(conf: Config) -> crate::Result<()> {
let mut buf = [0_u8; 65600];
match conf.command {
Op::Send { host, port } => {
let stream = UdpSocket::bind("0.0.0.0:0")?;
let target = format!("{}:{}",host,port).to_socket_addrs()?.next().unwrap();
let mut stdin = stdin().lock();
while let Ok(n) = stdin.read(&mut buf) {
if n == 0 { break }
let mut sent = 0;
loop {
sent += stream.send_to(&buf[sent..n], target)?;
if sent == n { break }
}
}
}
Op::Listen { port } => {
let ip = IpAddr::V6([0_u8; 16].into());
let socket = SocketAddr::new(ip, port);
let listener = UdpSocket::bind(socket)?;
let mut stdout = stdout().lock();
let (_,addr) = listener.peek_from(&mut buf)?;
listener.connect(addr)?;
while let Ok(n) = listener.recv(&mut buf) {
if n == 0 { break }
stdout.write_all(&buf[..n])?;
}
}
}
Ok(())
}
pub fn run() -> crate::Result<()> {
let conf = Config::parse();
if conf.udp {
udp(conf)
} else {
tcp(conf)
}
}