// Copyright 2016 Pierre-Étienne Meunier
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

//! ```
//! extern crate thrussh_client;
//! extern crate env_logger;
//! use thrussh_client::*;
//! use std::io::Write;
//!
//! fn main() {
//!     env_logger::init().unwrap();
//!     let mut client = thrussh_client::SSHClient::new("localhost", 22).unwrap();
//!     let key = thrussh_client::load_secret_key("/home/pe/.ssh/id_ed25519").unwrap();
//!     client.session().set_auth_public_key("pe".to_string(), key);
//!     client.authenticate().unwrap();
//!
//!     struct C;
//!     impl thrussh_client::Handler for C {
//!         fn data(&mut self,
//!                 _: u32,
//!                 _: Option<u32>,
//!                 data: &[u8],
//!                 _: &mut Session)
//!                 -> Result<(), Error> {
//!             try!(std::io::stdout().write(data));
//!             Ok(())
//!         }
//!     }
//!     let mut c = C;
//!     let channel = client.session().channel_open_session().unwrap();
//!     client.wait_channel_open(&mut c, channel).unwrap();
//!     client.session().exec(channel, false, "ls -l");
//!     client.session().eof(channel);
//!     client.wait_channel_close(&mut c, channel).unwrap();
//! }
//! ```


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(())
    }
}