use dirs;
use error::Error;
use futures::future::Either;
use futures::Future;
use meta;
use rpassword;
use std;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use thrussh::client::{Authenticate, Connection, Handler};
use thrussh::Tcp;
use thrussh_keys;
use tokio::io::{AsyncRead, AsyncWrite};
pub enum AuthAttempt {
Agent(thrussh_keys::key::PublicKey),
Key(Arc<thrussh_keys::key::KeyPair>),
Password(String),
}
#[derive(Debug)]
enum AuthState {
Agent(KeyPath),
Key(KeyPath),
Password,
}
pub struct AuthAttempts {
state: AuthState,
local_repo_root: Option<PathBuf>,
server_name: String,
}
impl AuthAttempts {
pub fn new(server_name: String, local_repo_root: Option<PathBuf>, use_agent: bool) -> Self {
AuthAttempts {
state: if use_agent {
AuthState::Agent(KeyPath::first())
} else {
AuthState::Key(KeyPath::first())
},
local_repo_root,
server_name,
}
}
}
#[derive(Debug, Clone, Copy)]
enum KeyLocation {
Local,
Pijul,
Ssh,
}
impl KeyLocation {
fn next(&self) -> Option<Self> {
match *self {
KeyLocation::Local => Some(KeyLocation::Pijul),
KeyLocation::Pijul => Some(KeyLocation::Ssh),
KeyLocation::Ssh => None,
}
}
}
#[derive(Debug, Clone, Copy)]
enum KeyType {
Ed25519,
Rsa,
}
#[derive(Debug, Clone, Copy)]
struct KeyPath {
location: KeyLocation,
typ: KeyType,
}
impl KeyPath {
fn first() -> Self {
KeyPath {
location: KeyLocation::Local,
typ: KeyType::Ed25519,
}
}
fn next(&self) -> Option<KeyPath> {
match self.typ {
KeyType::Ed25519 => {
if let Some(location) = self.location.next() {
Some(KeyPath {
location,
typ: KeyType::Ed25519,
})
} else {
Some(KeyPath {
location: KeyLocation::Local,
typ: KeyType::Rsa,
})
}
}
KeyType::Rsa => {
if let Some(location) = self.location.next() {
Some(KeyPath {
location,
typ: KeyType::Rsa,
})
} else {
None
}
}
}
}
}
impl AuthAttempts {
fn key_dir(&self, key: &KeyPath) -> Option<PathBuf> {
match key.location {
KeyLocation::Local => self.local_repo_root.clone(),
KeyLocation::Pijul => meta::global_path().ok(),
KeyLocation::Ssh => {
if let Some(mut path) = dirs::home_dir() {
path.push(".ssh");
Some(path)
} else {
None
}
}
}
}
fn key(&self, key: &KeyPath) -> Option<PathBuf> {
self.key_dir(key).map(|mut p| {
p.push(match key.typ {
KeyType::Ed25519 => "id_ed25519",
KeyType::Rsa => "id_rsa",
});
p
})
}
fn public_key(&self, key: &KeyPath) -> Option<PathBuf> {
self.key(key).map(|mut p| {
p.set_extension("pub");
p
})
}
}
impl Iterator for AuthAttempts {
type Item = AuthAttempt;
fn next(&mut self) -> Option<Self::Item> {
loop {
debug!("state {:?}", self.state);
match self.state {
AuthState::Agent(key_path) => {
let path = self.public_key(&key_path);
debug!("agent path {:?}", path);
if let Some(key_path) = key_path.next() {
self.state = AuthState::Agent(key_path)
} else {
self.state = AuthState::Key(KeyPath::first())
}
if let Some(path) = path {
if let Ok(key) = thrussh_keys::load_public_key(&path) {
return Some(AuthAttempt::Agent(key));
}
}
}
AuthState::Key(key_path) => {
let path = self.key(&key_path);
debug!("path {:?}", path);
if let Some(path) = path {
if let Some(key_path) = key_path.next() {
self.state = AuthState::Key(key_path)
} else {
self.state = AuthState::Password
}
if let Ok(key) = load_key_or_ask(&path) {
return Some(AuthAttempt::Key(Arc::new(key)));
}
} else {
self.state = AuthState::Password
}
}
AuthState::Password => {
let password = rpassword::prompt_password_stdout(&format!(
"Password for {:?}: ",
self.server_name
));
if let Ok(password) = password {
return Some(AuthAttempt::Password(password));
}
}
}
}
}
}
pub fn auth_attempt_future<R: Tcp + AsyncRead + AsyncWrite>(
co: Connection<R, super::remote::Client>,
auth_attempts: AuthAttempts,
user: String,
add_to_agent: thrussh_config::AddKeysToAgent,
) -> impl Future<Item = Connection<R, super::remote::Client>, Error = Error> {
super::fold_until::new(
futures::stream::iter_ok(auth_attempts),
co,
move |mut co, attempt| {
if let AuthAttempt::Key(key) = attempt {
debug!("not authenticated");
let agent_constraints = match add_to_agent {
thrussh_config::AddKeysToAgent::Yes => Some(&[][..]),
thrussh_config::AddKeysToAgent::No => None,
thrussh_config::AddKeysToAgent::Confirm => {
Some(&[thrussh_keys::agent::Constraint::Confirm][..])
}
thrussh_config::AddKeysToAgent::Ask => None, };
if let Some(cons) = agent_constraints {
if let Some(agent) = co.handler_mut().agent.take() {
let user = user.clone();
return Either::A(agent.add_identity(&key, cons).from_err().and_then(
move |(agent, _)| {
co.handler_mut().agent = Some(agent);
next_auth(co, &user, AuthAttempt::Key(key))
},
));
}
}
Either::B(next_auth(co, &user, AuthAttempt::Key(key)))
} else {
Either::B(next_auth(co, &user, attempt))
}
},
|co| futures::finished::<_, Error>((!co.is_authenticated(), co)),
)
}
fn next_auth<R: AsyncRead + AsyncWrite + Tcp, H: Handler>(
session: Connection<R, H>,
user: &str,
next: AuthAttempt,
) -> Authenticate<R, H> {
debug!("next_auth");
match next {
AuthAttempt::Agent(pk) => {
debug!("agent");
session.authenticate_key_future(user, pk)
}
AuthAttempt::Key(k) => {
debug!("key");
session.authenticate_key(user, k)
}
AuthAttempt::Password(pass) => {
debug!("password");
session.authenticate_password(user, pass)
}
}
}
pub fn load_key_or_ask(path_sec: &Path) -> Result<thrussh_keys::key::KeyPair, Error> {
debug!("path_sec {:?}", path_sec);
match thrussh_keys::load_secret_key(path_sec.to_str().unwrap(), None) {
Ok(key) => Ok(key),
Err(e) => match e {
thrussh_keys::Error::KeyIsEncrypted => {
let password = rpassword::prompt_password_stdout(&format!(
"Password for key {:?}: ",
path_sec
))?;
if password.is_empty() {
return Err(Error::EmptyPassword);
}
let key = thrussh_keys::load_secret_key(
path_sec.to_str().unwrap(),
Some(password.as_bytes()),
)?;
Ok(key)
}
thrussh_keys::Error::IO(ref e) if e.kind() == std::io::ErrorKind::NotFound => {
Err(Error::SshKeyNotFound {
path: path_sec.to_path_buf(),
})
}
_ => Err(From::from(e)),
},
}
}