Skip to content

Instantly share code, notes, and snippets.

@Freaky
Last active March 26, 2019 18:31
Show Gist options
  • Save Freaky/ab948a13d9438201e1876497559246d4 to your computer and use it in GitHub Desktop.
Save Freaky/ab948a13d9438201e1876497559246d4 to your computer and use it in GitHub Desktop.
Rust ssh tarpit
use std::env;
use std::net::SocketAddr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::time::{Duration, Instant};
use exitcode;
use env_logger;
use log::{error, info, warn};
use futures::future::{Loop, loop_fn};
use futures::stream::Stream;
use futures::Future;
use tokio::net::TcpListener;
use tokio::timer::Delay;
use rand::{thread_rng, Rng};
static NUM_CLIENTS: AtomicUsize = AtomicUsize::new(0);
#[cfg(feature="capsicum")]
use capsicum;
fn errx<M: AsRef<str>>(code: i32, message: M) {
error!("{}", message.as_ref());
std::process::exit(code);
}
fn main() {
env_logger::init();
let addr = env::args()
.nth(1)
.unwrap_or_else(|| "0.0.0.0:2222".to_string())
.parse::<SocketAddr>()
.map_err(|_| errx(exitcode::USAGE, "Error parsing listen address"))
.expect("unreachable");
let listener = TcpListener::bind(&addr)
.map_err(|err| errx(exitcode::OSERR, format!("bind(), error: {}", err)))
.expect("unreachable");
info!("listen, addr: {}", addr);
#[cfg(feature="capsicum")]
{
let _ = capsicum::enter();
info!("capsicum sandbox, enabled: {}", capsicum::sandboxed());
}
let server = listener
.incoming()
.map_err(|err| error!("accept(), error: {}", err))
.filter_map(|sock| {
sock.peer_addr()
.map_err(|err| error!("peer_addr(), error: {}", err))
.map(|peer| (sock, peer))
.ok()
})
.for_each(|(sock, peer)| {
let connected = NUM_CLIENTS.fetch_add(1, Ordering::Relaxed);
info!("connect, peer: {}, clients: {}", peer, connected + 1);
let start = Instant::now();
let _ = sock
.set_recv_buffer_size(1)
.map_err(|err| warn!("set_recv_buffer_size(), error: {}", err));
let tarpit = loop_fn(sock, move |sock| {
Delay::new(Instant::now() + Duration::from_secs(10))
.map_err(|err| {
error!("tokio timer, error: {}", err);
std::io::Error::new(std::io::ErrorKind::Other, "timer failure")
})
.and_then(move |_| {
tokio::io::write_all(sock, format!("{:x}\r\n", thread_rng().gen::<u32>()))
})
.map(|(sock, _)| Loop::Continue(sock))
.or_else(move |err| {
let connected = NUM_CLIENTS.fetch_sub(1, Ordering::Relaxed);
info!(
"disconnect, peer: {}, duration: {:.2?}, error: {}, clients: {}",
peer,
start.elapsed(),
err,
connected - 1
);
Ok(Loop::Break(()))
})
});
tokio::spawn(tarpit)
});
tokio::run(server);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment