extern crate thrussh;
extern crate ring;
extern crate crypto;
#[macro_use]
extern crate log;
extern crate base64;
extern crate untrusted;
extern crate byteorder;


use base64::{decode_config, encode_config, MIME};
use thrussh::{key, parse_public_key, Named};
use std::path::Path;
use std::borrow::Cow;

use std::fs::{File, OpenOptions};
use std::io::{BufReader, BufRead, Read, Write, Seek, SeekFrom};
use thrussh::encoding::Reader;
use ring::signature;
use byteorder::{BigEndian, WriteBytesExt};

#[derive(Debug)]
pub enum Error {
    CouldNotReadKey,
    IO(std::io::Error),
    RustCrypto(crypto::symmetriccipher::SymmetricCipherError),
    Ring(ring::error::Unspecified),
    Utf8(std::str::Utf8Error),
    Base64(base64::DecodeError),
    Thrussh(thrussh::Error),
    NoHomeDir,
    KeyChanged(usize)
}

impl From<std::io::Error> for Error {
    fn from(e: std::io::Error) -> Error {
        Error::IO(e)
    }
}
impl From<thrussh::Error> for Error {
    fn from(e: thrussh::Error) -> Error {
        Error::Thrussh(e)
    }
}
impl From<std::str::Utf8Error> for Error {
    fn from(e: std::str::Utf8Error) -> Error {
        Error::Utf8(e)
    }
}
impl From<base64::DecodeError> for Error {
    fn from(e: base64::DecodeError) -> Error {
        Error::Base64(e)
    }
}
impl From<crypto::symmetriccipher::SymmetricCipherError> for Error {
    fn from(e: crypto::symmetriccipher::SymmetricCipherError) -> Error {
        Error::RustCrypto(e)
    }
}
impl From<ring::error::Unspecified> for Error {
    fn from(e: ring::error::Unspecified) -> Error {
        Error::Ring(e)
    }
}


const KEYTYPE_ED25519: &'static [u8] = b"ssh-ed25519";
const KEYTYPE_RSA: &'static [u8] = b"ssh-rsa";

/// Load a public key from a file. Ed25519 and RSA keys are supported.
///
/// ```
/// thrussh_keys::load_public_key("/home/pe/.ssh/id_ed25519.pub").unwrap();
/// ```
pub fn load_public_key<P:AsRef<Path>>(path: P) -> Result<key::PublicKey, Error> {
    let mut pubkey = String::new();
    let mut file = try!(File::open(path.as_ref()));
    try!(file.read_to_string(&mut pubkey));

    debug!("decode_public_key: {:?}", pubkey);
    let mut split = pubkey.split_whitespace();
    match (split.next(), split.next()) {
        (Some(_), Some(key)) => parse_public_key_base64(key),
        (Some(key), None) => parse_public_key_base64(key),
        _ => Err(Error::CouldNotReadKey),
    }
}

/// Reads a public key from the standard encoding. In some cases, the
/// encoding is prefixed with a key type identifier and a space (such
/// as `ssh-ed25519 AAAAC3N...`).
///
/// ```
/// thrussh_keys::parse_public_key_base64("AAAAC3NzaC1lZDI1NTE5AAAAIJdD7y3aLq454yWBdwLWbieU1ebz9/cu7/QEXn9OIeZJ").is_ok();
/// ```
pub fn parse_public_key_base64(key: &str) -> Result<key::PublicKey, Error> {
    let base = decode_config(key, MIME)?;
    debug!("parse_public_key_base64: {:?}", base);
    Ok(parse_public_key(&base)?)
}


/// Write a public key onto the provided `Write`, encoded in base-64.
pub fn write_public_key_base64<W:Write>(mut w:W, publickey:&key::PublicKey) -> Result<(), Error> {
    try!(w.write_all(publickey.name().as_bytes()));
    try!(w.write_all(b" "));
    let mut s = Vec::new();
    let name = publickey.name().as_bytes();
    s.write_u32::<BigEndian>(name.len() as u32).unwrap();
    s.extend(name);
    match *publickey {
        key::PublicKey::Ed25519(ref publickey) => {
            s.write_u32::<BigEndian>(publickey.len() as u32).unwrap();
            s.extend(publickey);
        }
        _ => unimplemented!()
    }
    w.write_all(encode_config(&s, MIME).as_bytes())?;
    Ok(())
}

enum Format {
    Rsa,
    Openssh
}

/// Parse a (possibly password-protected) secret key. Only ED25519 and RSA keys are currently supported.
///
/// ```
/// let secret = "-----BEGIN OPENSSH PRIVATE KEY-----
/// b3BlbnNzaC1rZXktdjEAAAAACmFlczI1Ni1jYmMAAAAGYmNyeXB0AAAAGAAAABDLGyfA39
/// J2FcJygtYqi5ISAAAAEAAAAAEAAAAzAAAAC3NzaC1lZDI1NTE5AAAAIN+Wjn4+4Fcvl2Jl
/// KpggT+wCRxpSvtqqpVrQrKN1/A22AAAAkOHDLnYZvYS6H9Q3S3Nk4ri3R2jAZlQlBbUos5
/// FkHpYgNw65KCWCTXtP7ye2czMC3zjn2r98pJLobsLYQgRiHIv/CUdAdsqbvMPECB+wl/UQ
/// e+JpiSq66Z6GIt0801skPh20jxOO3F52SoX1IeO5D5PXfZrfSZlw6S8c7bwyp2FHxDewRx
/// 7/wNsnDM0T7nLv/Q==
/// -----END OPENSSH PRIVATE KEY-----
/// ";
/// thrussh_keys::load_secret_key(secret, Some(b"blabla"), None).unwrap();
/// ```
pub fn load_secret_key(secret_: &str, password: Option<&[u8]>, public: Option<&str>) -> Result<key::Algorithm, Error> {

    let mut secret = String::new();
    let mut started = false;

    let mut format = None;
    for l in secret_.lines() {
        if l == "-----BEGIN OPENSSH PRIVATE KEY-----" {
            format = Some(Format::Openssh);
            started = true
        } else if l == "-----BEGIN RSA PRIVATE KEY-----" {
            format = Some(Format::Rsa);
            started = true
        } else if l == "-----END OPENSSH PRIVATE KEY-----" || l == "-----END RSA PRIVATE KEY-----" {
            break;
        } else if started {
            secret.push_str(&l)
        }
    }
    debug!("secret: {:?}", secret);
    let secret = decode_config(&secret, MIME)?;
    match format {
        Some(Format::Openssh) => {
            debug!("secret: {:?}", secret);
            if &secret[0..15] == b"openssh-key-v1\0" {
                let mut position = secret.reader(15);

                let ciphername = try!(position.read_string());
                let kdfname = try!(position.read_string());
                let kdfoptions = try!(position.read_string());
                info!("ciphername: {:?}", std::str::from_utf8(ciphername));
                debug!("kdf: {:?} {:?}",
                       std::str::from_utf8(kdfname),
                       kdfoptions);

                let nkeys = try!(position.read_u32());

                for _ in 0..nkeys {
                    let public_string = try!(position.read_string());
                    let mut pos = public_string.reader(0);
                    let t = try!(pos.read_string());
                    let pubkey = pos.read_string()?;
                    if t == KEYTYPE_ED25519 {
                        info!("public: ED25519:{:?}", pubkey);
                    } else if t == KEYTYPE_RSA {
                        info!("public: RSA:{:?}", pubkey);
                    } else {
                        info!("warning: no public key");
                    }
                }
                info!("there are {} keys in this file", nkeys);
                let secret_ = try!(position.read_string());
                let mut secret = secret_.to_vec();
                decrypt_secret_key(
                    ciphername, kdfname, kdfoptions, password, secret_, &mut secret
                )?;
                let mut position = secret.reader(0);
                let check0 = try!(position.read_u32());
                let check1 = try!(position.read_u32());
                debug!("check0: {:?}", check0);
                debug!("check1: {:?}", check1);
                for _ in 0..nkeys {

                    let key_type = try!(position.read_string());
                    let pubkey = try!(position.read_string());
                    debug!("pubkey = {:?}", pubkey);
                    let seckey = try!(position.read_string());
                    let comment = try!(position.read_string());
                    debug!("comment = {:?}", comment);

                    if key_type == KEYTYPE_ED25519 {
                        let (a,b) = seckey.split_at(32);
                        assert_eq!(pubkey, b);
                        let keypair = try!(signature::Ed25519KeyPair::from_seed_and_public_key(untrusted::Input::from(a), untrusted::Input::from(pubkey)));
                        let keypair = key::Algorithm::Ed25519(keypair);
                        return Ok(keypair)
                    } else {
                        info!("unsupported key type {:?}", std::str::from_utf8(key_type));
                    }
                }
                Err(Error::CouldNotReadKey)

            } else {
                Err(Error::CouldNotReadKey)
            }
        }
        Some(Format::Rsa) => {
            debug!("secret: {:?}", secret);
            let keypair = signature::RSAKeyPair::from_der(untrusted::Input::from(&secret))?;
            if let Some(public) = public {
                let mut split = public.split_whitespace();
                if let (Some("ssh-rsa"), Some(public)) = (split.next(), split.next()) {
                    if let key::PublicKey::RSA(public) = parse_public_key(public.as_bytes())? {
                        return Ok(key::Algorithm::RSA(
                            std::sync::Arc::new(keypair),
                            public,
                        ))
                    }
                }
            }
            Err(Error::CouldNotReadKey)
        }
        None => Err(Error::CouldNotReadKey)
    }
}

fn decrypt_secret_key(
    ciphername: &[u8],
    kdfname: &[u8], kdfoptions: &[u8], password: Option<&[u8]>,
    secret_key: &[u8], output: &mut [u8]
) -> Result<(), Error> {

    if kdfname == b"none" {
        if password.is_none() {
            debug!("key: {:?}", secret_key);
            output.clone_from_slice(secret_key);
            return Ok(())
        }
    } else if let Some(password) = password {

        let mut key =
            match ciphername {
                b"aes256-cbc" => vec![0; 16 + 32],
                _ => return Err(Error::CouldNotReadKey)
            };

        match kdfname {
            b"bcrypt" => {
                let mut kdfopts = kdfoptions.reader(0);
                let salt = kdfopts.read_string()?;
                let rounds = kdfopts.read_u32()?;
                debug!("salt = {:?}, rounds = {:?}", salt, rounds);
                crypto::bcrypt_pbkdf::bcrypt_pbkdf(password, salt, rounds, &mut key);
            },
            _ => return Err(Error::CouldNotReadKey)
        };
        debug!("key = {:?}", key);
        let (key, iv) = key.split_at(32);
        let mut refread = crypto::buffer::RefReadBuffer::new(secret_key);
        let mut refwrite = crypto::buffer::RefWriteBuffer::new(output);
        crypto::aes::cbc_decryptor(
            crypto::aes::KeySize::KeySize256,
            key, iv, crypto::blockmodes::NoPadding
        ).decrypt(&mut refread, &mut refwrite, true)?;
        return Ok(())
    }
    info!("unsupported secret key cipher: {:?}",
          std::str::from_utf8(kdfname));
    Err(Error::CouldNotReadKey)
}



/// Record a host's public key into a nonstandard location.
pub fn learn_known_hosts_path<P:AsRef<Path>>(host:&str, port:u16, pubkey:&key::PublicKey, path:P) -> Result<(), Error> {


    let mut file = try!(OpenOptions::new()
                        .read(true)
                        .append(true)
                        .create(true)
                        .open(path));

    // Test whether the known_hosts file ends with a \n
    let mut buf = [0;1];
    try!(file.seek(SeekFrom::End(-1)));
    try!(file.read_exact(&mut buf));
    let ends_in_newline = buf[0] == b'\n';

    // Write the key.
    try!(file.seek(SeekFrom::Start(0)));
    let mut file = std::io::BufWriter::new(file);
    if !ends_in_newline {
        try!(write!(file, "\n"));
    }
    if port != 22 {
        try!(write!(file, "[{}]:{} ", host, port))
    } else {
        try!(write!(file, "{} ", host))
    }
    try!(write_public_key_base64(&mut file, pubkey));
    try!(write!(file, "\n"));
    Ok(())
}

/// Check that a server key matches the one recorded in file `path`.
pub fn check_known_hosts_path<P: AsRef<Path>>(host: &str,
                                              port: u16,
                                              pubkey: &key::PublicKey,
                                              path: P)
                                              -> Result<bool, Error> {
    let mut f = BufReader::new(try!(File::open(path)));
    let mut buffer = String::new();

    let host_port = if port == 22 {
        Cow::Borrowed(host)
    } else {
        Cow::Owned(format!("[{}]:{}", host, port))
    };
    let mut line = 1;
    while f.read_line(&mut buffer).unwrap() > 0 {
        debug!("check_known_hosts_path: {:?}", buffer);
        {
            if buffer.as_bytes()[0] == b'#' {
                buffer.clear();
                continue;
            }
            let mut s = buffer.split(' ');
            let hosts = s.next();
            let _ = s.next();
            let key = s.next();
            debug!("{:?} {:?}", hosts, key);
            match (hosts, key) {
                (Some(h), Some(k)) => {
                    let host_matches = h.split(',').any(|x| x == host_port);
                    debug!("host matches: {:?}", host_matches);
                    // debug!("{:?} {:?}", parse_public_key_base64(k), pubkey);
                    if host_matches {
                        if &try!(parse_public_key_base64(k)) == pubkey {
                            return Ok(true);
                        } else {
                            return Err(Error::KeyChanged(line));
                        }
                    }

                }
                _ => {}
            }
        }
        buffer.clear();
        line += 1;
    }
    Ok(false)
}


/// Record a host's public key into the user's known_hosts file.
#[cfg(target_os = "windows")]
pub fn learn_known_hosts(host: &str, port: u16, pubkey: &key::PublicKey) -> Result<(), Error> {
    if let Some(mut known_host_file) = std::env::home_dir() {
        known_host_file.push("ssh");
        known_host_file.push("known_hosts");
        learn_known_hosts_path(host, port, pubkey, &known_host_file)
    } else {
        Err(Error::NoHomeDir)
    }
}

/// Record a host's public key into the user's known_hosts file.
#[cfg(not(target_os = "windows"))]
pub fn learn_known_hosts(host: &str, port: u16, pubkey: &key::PublicKey) -> Result<(), Error> {
    if let Some(mut known_host_file) = std::env::home_dir() {
        known_host_file.push(".ssh");
        known_host_file.push("known_hosts");
        learn_known_hosts_path(host, port, pubkey, &known_host_file)
    } else {
        Err(Error::NoHomeDir)
    }
}

/// Check whether the host is known, from its standard location.
#[cfg(target_os = "windows")]
pub fn check_known_hosts(host: &str, port: u16, pubkey: &key::PublicKey) -> Result<bool, Error> {
    if let Some(mut known_host_file) = std::env::home_dir() {
        known_host_file.push("ssh");
        known_host_file.push("known_hosts");
        check_known_hosts_path(host, port, pubkey, &known_host_file)
    } else {
        Err(Error::NoHomeDir)
    }
}

/// Check whether the host is known, from its standard location.
#[cfg(not(target_os = "windows"))]
pub fn check_known_hosts(host: &str, port: u16, pubkey: &key::PublicKey) -> Result<bool, Error> {
    if let Some(mut known_host_file) = std::env::home_dir() {
        known_host_file.push(".ssh");
        known_host_file.push("known_hosts");
        debug!("known_hosts file = {:?}", known_host_file);
        check_known_hosts_path(host, port, pubkey, &known_host_file)
    } else {
        Err(Error::NoHomeDir)
    }
}


#[cfg(test)]
mod test {
    extern crate tempdir;
    extern crate env_logger;
    use std::fs::File;
    use std::io::Write;
    use super::*;

    #[test]
    fn test_check_known_hosts() {
        env_logger::init().unwrap_or(());
        let dir = tempdir::TempDir::new("thrussh").unwrap();
        let path = dir.path().join("known_hosts");
        {
            let mut f = File::create(&path).unwrap();
            f.write(b"[localhost]:13265 ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIJdD7y3aLq454yWBdwLWbieU1ebz9/cu7/QEXn9OIeZJ\n#pijul.org,37.120.161.53 ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIA6rWI3G2sz07DnfFlrouTcysQlj2P+jpNSOEWD9OJ3X\npijul.org,37.120.161.53 ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIA6rWI3G1sz07DnfFlrouTcysQlj2P+jpNSOEWD9OJ3X\n").unwrap();
        }

        // Valid key, non-standard port.
        let host = "localhost";
        let port = 13265;
        let hostkey = parse_public_key_base64("AAAAC3NzaC1lZDI1NTE5AAAAIJdD7y3aLq454yWBdwLWbieU1ebz9/cu7/QEXn9OIeZJ")
            .unwrap();
        assert!(check_known_hosts_path(host, port, &hostkey, &path).unwrap());

        // Valid key, several hosts, port 22
        let host = "pijul.org";
        let port = 22;
        let hostkey = parse_public_key_base64("AAAAC3NzaC1lZDI1NTE5AAAAIA6rWI3G1sz07DnfFlrouTcysQlj2P+jpNSOEWD9OJ3X")
            .unwrap();
        assert!(check_known_hosts_path(host, port, &hostkey, &path).unwrap());

        // Now with the key in a comment above, check that it's not recognized
        let host = "pijul.org";
        let port = 22;
        let hostkey = parse_public_key_base64("AAAAC3NzaC1lZDI1NTE5AAAAIA6rWI3G2sz07DnfFlrouTcysQlj2P+jpNSOEWD9OJ3X")
            .unwrap();
        assert!(check_known_hosts_path(host, port, &hostkey, &path).is_err());
    }
}