extern crate mio;
extern crate thrussh;
use mio::{Token, EventSet, Poll, PollOpt};
use mio::tcp::TcpStream;
#[macro_use]
extern crate log;
use std::io::{Read, Write, BufReader};
pub use thrussh::{Error, load_secret_key};
pub use thrussh::client::{Handler, Session};
use thrussh::{CryptoBuf, key};
use std::net::ToSocketAddrs;
use std::default::Default;
use std::sync::Arc;
#[derive(Debug)]
enum RunUntil {
Authenticated,
ChannelOpened(u32),
ChannelClosed(u32),
}
pub struct SSHClient<'b> {
poll: Poll,
host: &'b str,
port: u16,
buffer0: CryptoBuf,
buffer1: CryptoBuf,
pub connection: thrussh::client::Connection,
stream: BufReader<TcpStream>,
}
struct C<'b> {
host: &'b str,
port: u16,
}
impl<'b> thrussh::client::Handler for C<'b> {
fn check_server_key(&mut self, pubkey: &key::PublicKey) -> Result<bool, thrussh::Error> {
thrussh::check_known_hosts(self.host, self.port, pubkey)
}
}
pub struct Basic;
impl thrussh::client::Handler for Basic {}
impl<'b> SSHClient<'b> {
pub fn new(host: &'b str, port: u16) -> Result<Self, Error> {
let addr = try!((host, port).to_socket_addrs()).next().unwrap();
let sock = try!(TcpStream::connect(&addr));
let mut poll = Poll::new().unwrap();
try!(poll.register(&sock, Token(0), EventSet::all(), PollOpt::edge()));
Ok(SSHClient {
poll: poll,
host: host,
port: port,
buffer0: CryptoBuf::new(),
buffer1: CryptoBuf::new(),
stream: BufReader::new(sock),
connection: thrussh::client::Connection::new(Arc::new(Default::default())),
})
}
pub fn authenticate(&mut self) -> Result<bool, Error> {
try!(self.poll
.reregister(self.stream.get_ref(),
Token(0),
EventSet::all(),
PollOpt::edge()));
let mut d = C {
host: self.host,
port: self.port,
};
try!(self.run(&mut d, Some(RunUntil::Authenticated)));
Ok(self.connection.session.is_authenticated())
}
pub fn wait_channel_open<C: thrussh::client::Handler>(&mut self,
c: &mut C,
channel: u32)
-> Result<(), Error> {
try!(self.poll
.reregister(self.stream.get_ref(),
Token(0),
EventSet::all(),
PollOpt::edge()));
try!(self.run(c, Some(RunUntil::ChannelOpened(channel))));
Ok(())
}
pub fn wait_channel_close<C: thrussh::client::Handler>(&mut self,
c: &mut C,
channel: u32)
-> Result<(), Error> {
try!(self.poll
.reregister(self.stream.get_ref(),
Token(0),
EventSet::all(),
PollOpt::edge()));
try!(self.run(c, Some(RunUntil::ChannelClosed(channel))));
Ok(())
}
pub fn session(&mut self) -> &mut Session {
&mut self.connection.session
}
fn run<R: thrussh::client::Handler>(&mut self,
client: &mut R,
until: Option<RunUntil>)
-> Result<(), Error> {
try!(self.connection.write(self.stream.get_mut()));
loop {
match self.poll.poll(None) {
Ok(n) if n > 0 => {
let events = self.poll.event(0).kind;
if events.is_error() || events.is_hup() {
return Err(Error::HUP);
} else {
if events.is_readable() {
try!(self.connection.read(client,
&mut self.stream,
&mut self.buffer0,
&mut self.buffer1));
match until {
Some(RunUntil::Authenticated) if self.connection
.session
.is_authenticated() => {
return Ok(());
}
Some(RunUntil::Authenticated) if self.connection
.session
.needs_auth_method() => {
return Ok(());
}
Some(RunUntil::ChannelOpened(x)) if self.connection
.session
.channel_is_open(x) => {
return Ok(());
}
Some(RunUntil::ChannelClosed(x)) if !self.connection
.session
.channel_is_open(x) => {
return Ok(());
}
_ => {}
}
}
if events.is_writable() {
try!(self.connection.write(self.stream.get_mut()));
}
}
}
_ => break,
}
}
Ok(())
}
pub fn run_until<R: thrussh::client::Handler, F: Fn(&mut R) -> bool>(&mut self,
client: &mut R,
until: F)
-> Result<(), Error> {
try!(self.connection.write(self.stream.get_mut()));
while !until(client) {
match self.poll.poll(None) {
Ok(n) if n > 0 => {
let events = self.poll.event(0).kind;
if events.is_error() || events.is_hup() {
return Err(Error::HUP);
} else {
if events.is_readable() {
try!(self.connection.read(client,
&mut self.stream,
&mut self.buffer0,
&mut self.buffer1));
}
if events.is_writable() {
try!(self.connection.write(self.stream.get_mut()));
}
}
}
_ => break,
}
}
Ok(())
}
}