diff --git a/crates/keyfork-shard/src/lib.rs b/crates/keyfork-shard/src/lib.rs index 5f94be5..bb3a59d 100644 --- a/crates/keyfork-shard/src/lib.rs +++ b/crates/keyfork-shard/src/lib.rs @@ -2,9 +2,10 @@ #![allow(clippy::expect_fun_call)] use std::{ - io::{stdin, stdout, Read, Write}, + io::{Read, Write}, rc::Rc, - sync::Mutex, + str::FromStr, + sync::{LazyLock, Mutex}, }; use aes_gcm::{ @@ -22,7 +23,7 @@ use keyfork_prompt::{ mnemonic::{MnemonicSetValidator, MnemonicValidator, WordLength}, Validator, }, - Message as PromptMessage, PromptHandler, Terminal, + Message as PromptMessage, PromptHandler, }; use sha2::Sha256; use x25519_dalek::{EphemeralSecret, PublicKey}; @@ -34,6 +35,30 @@ const PLAINTEXT_LENGTH: u8 = 32 // shard + 1; // length; const ENCRYPTED_LENGTH: u8 = PLAINTEXT_LENGTH + 16; +#[derive(PartialEq, Eq, Clone, Copy)] +enum RetryScanMnemonic { + Retry, + Continue, +} + +impl keyfork_prompt::Choice for RetryScanMnemonic { + fn identifier(&self) -> Option { + Some(match self { + RetryScanMnemonic::Retry => 'r', + RetryScanMnemonic::Continue => 'c', + }) + } +} + +impl std::fmt::Display for RetryScanMnemonic { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + RetryScanMnemonic::Retry => write!(f, "Retry scanning mnemonic."), + RetryScanMnemonic::Continue => write!(f, "Continue to manual mnemonic entry."), + } + } +} + #[cfg(feature = "openpgp")] pub mod openpgp; @@ -247,19 +272,28 @@ pub trait Format { .lock() .expect(bug!(POISONED_MUTEX)) .prompt_message(PromptMessage::Text(QRCODE_PROMPT.to_string()))?; - if let Ok(Some(qrcode_content)) = - keyfork_qrcode::scan_camera(std::time::Duration::from_secs(30), 0) - { - let decoded_data = BASE64_STANDARD - .decode(qrcode_content) - .expect(bug!("qrcode should contain base64 encoded data")); - pubkey_data = Some(decoded_data.try_into().map_err(|_| InvalidData)?) - } else { - prompt - .lock() - .expect(bug!(POISONED_MUTEX)) - .prompt_message(PromptMessage::Text(QRCODE_ERROR.to_string()))?; - }; + loop { + if let Ok(Some(qrcode_content)) = keyfork_qrcode::scan_camera( + std::time::Duration::from_secs(*QRCODE_TIMEOUT), + 0, + ) { + let decoded_data = BASE64_STANDARD + .decode(qrcode_content) + .expect(bug!("qrcode should contain base64 encoded data")); + pubkey_data = Some(decoded_data.try_into().map_err(|_| InvalidData)?); + break; + } else { + let mut prompt = prompt.lock().expect(bug!(POISONED_MUTEX)); + let choice = keyfork_prompt::prompt_choice( + &mut **prompt, + "A QR code could not be scanned. Retry or continue?", + &[RetryScanMnemonic::Retry, RetryScanMnemonic::Continue], + )?; + if choice == RetryScanMnemonic::Continue { + break; + } + }; + } } // if QR code scanning failed or was unavailable, read from a set of mnemonics @@ -459,9 +493,13 @@ pub(crate) const HUNK_VERSION: u8 = 2; pub(crate) const HUNK_OFFSET: usize = 2; const QRCODE_PROMPT: &str = "Press enter, then present QR code to camera."; -const QRCODE_TIMEOUT: u64 = 60; // One minute const QRCODE_COULDNT_READ: &str = "A QR code could not be scanned. Please enter their words: "; -const QRCODE_ERROR: &str = "Unable to scan a QR code. Falling back to text entry."; +static QRCODE_TIMEOUT: LazyLock = LazyLock::new(|| { + std::env::var("KEYFORK_QRCODE_TIMEOUT") + .ok() + .and_then(|t| u64::from_str(&t).ok()) + .unwrap_or(60) +}); /// Establish ECDH transport for remote operators, receive transport-encrypted shares, decrypt the /// shares, and combine them. @@ -476,7 +514,7 @@ const QRCODE_ERROR: &str = "Unable to scan a QR code. Falling back to text entry /// The function may panic if it is given payloads generated using a version of Keyfork that is /// incompatible with the currently running version. pub fn remote_decrypt(w: &mut impl Write) -> Result<(), Box> { - let mut pm = Terminal::new(stdin(), stdout())?; + let mut pm = keyfork_prompt::default_handler()?; let mut iter_count = None; let mut shares = vec![]; @@ -523,23 +561,34 @@ pub fn remote_decrypt(w: &mut impl Write) -> Result<(), Box Result<(), Box( - &mut pm, + &mut *pm, QRCODE_COULDNT_READ, 3, &*validator.to_fn(), diff --git a/crates/util/keyfork-bug/src/lib.rs b/crates/util/keyfork-bug/src/lib.rs index 7249480..926529a 100644 --- a/crates/util/keyfork-bug/src/lib.rs +++ b/crates/util/keyfork-bug/src/lib.rs @@ -16,6 +16,12 @@ //! ``` //! //! ```rust,should_panic +//! let rows = 24; +//! let input_lines_len = 25; +//! assert!(input_lines_len < rows, "{input_lines_len} can't fit in {rows} lines!"); +//! ``` +//! +//! ```rust,should_panic //! use std::fs::File; //! use keyfork_bug as bug; //! @@ -83,6 +89,29 @@ macro_rules! bug { }}; } +/// Assert a condition is true, otherwise throwing an error using Keyfork Bug. +/// +/// # Examples +/// ```rust +/// let expectations = "conceivable!"; +/// let circumstances = "otherwise"; +/// assert!(circumstances != expectations, "you keep using that word..."); +/// ``` +/// +/// Variables can be used in the error message, without having to pass them manually. +/// +/// ```rust,should_panic +/// let rows = 24; +/// let input_lines_len = 25; +/// assert!(input_lines_len < rows, "{input_lines_len} can't fit in {rows} lines!"); +/// ``` +#[macro_export] +macro_rules! assert { + ($cond:expr, $($input:tt)*) => { + std::assert!($cond, "{}", keyfork_bug::bug!($($input)*)); + } +} + /// Return a closure that, when called, panics with a bug report message for Keyfork. Returning a /// closure can help handle the `clippy::expect_fun_call` lint. The closure accepts an error /// argument, so it is suitable for being used with [`Result`] types instead of [`Option`] types. diff --git a/crates/util/keyfork-prompt/examples/test-basic-prompt.rs b/crates/util/keyfork-prompt/examples/test-basic-prompt.rs index 5aca0d2..cc453af 100644 --- a/crates/util/keyfork-prompt/examples/test-basic-prompt.rs +++ b/crates/util/keyfork-prompt/examples/test-basic-prompt.rs @@ -1,15 +1,47 @@ #![allow(missing_docs)] -use keyfork_prompt::{ - Message, - default_handler, -}; +use keyfork_prompt::default_handler; + +#[derive(PartialEq, Eq, Debug, Copy, Clone)] +pub enum Choices { + Retry, + Continue, +} + +impl std::fmt::Display for Choices { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Choices::Retry => write!( + f, + "Retry with some really long text that I want to cause issues with." + ), + Choices::Continue => write!( + f, + "Continue with some really long text that I want to cause issues with." + ), + } + } +} + +impl keyfork_prompt::Choice for Choices { + fn identifier(&self) -> Option { + Some(match self { + Choices::Retry => 'r', + Choices::Continue => 'c', + }) + } +} fn main() -> Result<(), Box> { let mut handler = default_handler()?; - let output = handler.prompt_input("Test message: ")?; - handler.prompt_message(Message::Text(format!("Result: {output}")))?; + let choice = keyfork_prompt::prompt_choice( + &mut *handler, + "Here are some options!", + &[Choices::Retry, Choices::Continue], + ); + + dbg!(&choice); Ok(()) } diff --git a/crates/util/keyfork-prompt/src/headless.rs b/crates/util/keyfork-prompt/src/headless.rs index 2f3ce01..c3287d2 100644 --- a/crates/util/keyfork-prompt/src/headless.rs +++ b/crates/util/keyfork-prompt/src/headless.rs @@ -4,9 +4,12 @@ //! directly intended to be machine-readable, but can be used for scriptable automation in a //! fashion similar to a terminal handler. -use std::io::{IsTerminal, Write}; +use std::{ + io::{IsTerminal, Write}, + str::FromStr, +}; -use crate::{BoxResult, Error, Message, PromptHandler, Result}; +use crate::{BoxResult, Choice, Error, Message, PromptHandler, Result}; /// A headless prompt handler, usable in situations when a terminal might not be available, or for /// scripting purposes where manual input from a terminal is not desirable. @@ -58,17 +61,47 @@ impl PromptHandler for Headless { fn prompt_message(&mut self, prompt: Message) -> Result<()> { match prompt { Message::Text(s) => { - self.stderr.write_all(s.as_bytes())?; + writeln!(&mut self.stderr, "{s}")?; self.stderr.flush()?; } Message::Data(s) => { - self.stderr.write_all(s.as_bytes())?; + writeln!(&mut self.stderr, "{s}")?; self.stderr.flush()?; } } + writeln!(&mut self.stderr, "Press enter to continue.")?; + self.stdin.read_line(&mut String::new())?; Ok(()) } + fn prompt_choice_num(&mut self, prompt: &str, choices: &[Box]) -> Result { + writeln!(&mut self.stderr, "{prompt}")?; + for (i, choice) in choices.iter().enumerate() { + match choice.identifier() { + Some(identifier) => { + writeln!(&mut self.stderr, "{i}. ({identifier})\t{choice}")?; + } + None => { + writeln!(&mut self.stderr, "{i}.\t{choice}")?; + } + } + } + self.stderr.flush()?; + let mut line = String::new(); + self.stdin.read_line(&mut line)?; + let selector_char = line.chars().next(); + if let Some(selector @ ('a'..='z' | 'A'..='Z')) = selector_char { + if let Some((index, _)) = choices.iter().enumerate().find(|(_, choice)| { + choice + .identifier() + .is_some_and(|identifier| selector == identifier) + }) { + return Ok(index); + } + } + usize::from_str(line.trim()).map_err(|e| Error::Custom(e.to_string())) + } + fn prompt_validated_wordlist( &mut self, prompt: &str, @@ -85,7 +118,7 @@ impl PromptHandler for Headless { self.stdin.read_line(&mut line)?; if let Err(e) = validator_fn(std::mem::take(&mut line)) { last_error = e.to_string(); - self.stderr.write_all(e.to_string().as_bytes())?; + writeln!(&mut self.stderr, "{e}")?; self.stderr.flush()?; } else { return Ok(()); @@ -108,8 +141,7 @@ impl PromptHandler for Headless { self.stdin.read_line(&mut line)?; if let Err(e) = validator_fn(std::mem::take(&mut line)) { last_error = e.to_string(); - self.stderr.write_all(e.to_string().as_bytes())?; - self.stderr.write_all(b"\n")?; + writeln!(&mut self.stderr, "{e}")?; self.stderr.flush()?; } else { return Ok(()); diff --git a/crates/util/keyfork-prompt/src/lib.rs b/crates/util/keyfork-prompt/src/lib.rs index 149b1c8..79a0ddd 100644 --- a/crates/util/keyfork-prompt/src/lib.rs +++ b/crates/util/keyfork-prompt/src/lib.rs @@ -50,6 +50,10 @@ pub enum Error { /// An error occurred while interacting with a terminal. #[error("IO Error: {0}")] IO(#[from] std::io::Error), + + /// An unexpected error occurred. + #[error("{0}")] + Custom(String), } #[allow(missing_docs)] @@ -64,6 +68,21 @@ pub enum Message { Data(String), } +/// A type that may represent an identifier to be used when using a choice prompt. +pub trait Choice: std::fmt::Display { + /// The identifier for the type. + fn identifier(&self) -> Option { + None + } +} + +// this way, we can make Box from &T +impl Choice for &T { + fn identifier(&self) -> Option { + Choice::identifier(*self) + } +} + #[doc(hidden)] pub type BoxResult = std::result::Result<(), Box>; @@ -98,6 +117,16 @@ pub trait PromptHandler { /// occurred while waiting for the user to dismiss the message. fn prompt_message(&mut self, prompt: Message) -> Result<()>; + /// Prompt the user for a choice between the provided options. The returned value is the index + /// of the given choice. + /// + /// This method SHOULD NOT be used directly. Instead, use [`prompt_choice`]. + /// + /// # Errors + /// The method may return an error if the message was not able to be displayed or if the input + /// could not be read. + fn prompt_choice_num(&mut self, prompt: &str, choices: &[Box]) -> Result; + /// Prompt the user for input based on a wordlist, while validating the wordlist using a /// provided parser function, returning the type from the parser. A language must be specified /// as the generic parameter `X` (any type implementing [`Wordlist`]) when parsing a wordlist. @@ -133,6 +162,29 @@ pub trait PromptHandler { ) -> Result<(), Error>; } +/// Prompt the user for a choice between the provided options. The returned value is the selected +/// choice. +/// +/// # Errors +/// The method may return an error if the message was not able to be displayed or if the input +/// could not be read. +#[allow(clippy::missing_panics_doc)] +pub fn prompt_choice( + handler: &mut dyn PromptHandler, + prompt: &str, + choices: &'static [T], +) -> Result +where + T: Choice + Copy + 'static, +{ + let boxed_choices = choices + .iter() + .map(|c| Box::new(c) as Box) + .collect::>(); + let choice = handler.prompt_choice_num(prompt, boxed_choices.as_slice())?; + Ok(choices[choice]) +} + /// Prompt the user for input based on a wordlist, while validating the wordlist using a /// provided parser function, returning the type from the parser. A language must be specified /// as the generic parameter `X` (any type implementing [`Wordlist`]) when parsing a wordlist. diff --git a/crates/util/keyfork-prompt/src/terminal.rs b/crates/util/keyfork-prompt/src/terminal.rs index 19c41bb..f4c85bc 100644 --- a/crates/util/keyfork-prompt/src/terminal.rs +++ b/crates/util/keyfork-prompt/src/terminal.rs @@ -21,7 +21,7 @@ use keyfork_crossterm::{ use keyfork_bug::bug; -use crate::{BoxResult, Error, Message, PromptHandler}; +use crate::{BoxResult, Choice, Error, Message, PromptHandler}; #[allow(missing_docs)] pub type Result = std::result::Result; @@ -129,21 +129,26 @@ where { fn drop(&mut self) { self.write + .execute(cursor::EnableBlinking) + .expect(bug!("can't enable blinking")) + .execute(cursor::Show) + .expect(bug!("can't show cursor")) .execute(DisableBracketedPaste) .expect(bug!("can't restore bracketed paste")); - self.write - .queue(terminal::Clear(terminal::ClearType::All)) - .expect(bug!("can't clear screen")) - .queue(cursor::MoveTo(0, 0)) - .expect(bug!("can't move to origin")) - .flush() - .expect(bug!("can't execute clear+move")); - self.write - .execute(LeaveAlternateScreen) - .expect(bug!("can't leave alternate screen")); self.terminal .disable_raw_mode() .expect(bug!("can't disable raw mode")); + // we don't want to clear error messages + if !std::thread::panicking() { + self.write + .queue(LeaveAlternateScreen) + .expect(bug!("can't leave alternate screen")) + .queue(terminal::Clear(terminal::ClearType::All)) + .expect(bug!("can't clear screen")) + .queue(cursor::MoveTo(0, 0)) + .expect(bug!("can't move to origin")); + } + self.write.flush().expect(bug!("can't execute terminal reset commands")); } } @@ -195,9 +200,7 @@ where prefix_length = line.len(); terminal.queue(Print(line))?; if lines.peek().is_some() { - terminal - .queue(cursor::MoveDown(1))? - .queue(cursor::MoveToColumn(0))?; + terminal.queue(cursor::MoveToNextLine(1))?; } } terminal.flush()?; @@ -264,6 +267,103 @@ where Ok(input) } + fn prompt_choice_num(&mut self, prompt: &str, choices: &[Box]) -> Result { + let mut terminal = self.lock().alternate_screen()?.raw_mode()?; + terminal + .queue(terminal::Clear(terminal::ClearType::All))? + .queue(cursor::MoveTo(0, 0))? + .queue(cursor::Hide)?; + + for line in prompt.lines() { + terminal + .queue(Print(line))? + .queue(cursor::MoveToNextLine(1))?; + + terminal.flush()?; + } + + let mut active_choice = 0; + let mut drawn = false; + + loop { + let (cols, rows) = terminal.size()?; + // all choices, plus their padding, plus the spacing between, minus whitespace at end. + let max_size = choices + .iter() + .fold(0usize, |agg, choice| agg + choice.to_string().len() + 2) + + std::cmp::max(choices.len(), 1) + - 1; + let horizontal = max_size < cols.into(); + keyfork_bug::assert!( + horizontal || usize::from(rows) > prompt.lines().count() + choices.len(), + "screen too small, can't fit choices on {rows}x{cols}", + ); + if horizontal { + terminal.queue(cursor::MoveToColumn(0))?; + } else if drawn { + terminal + .queue(cursor::MoveUp( + choices + .len() + .saturating_sub(1) + .try_into() + .expect(keyfork_bug::bug!("more than {} choices provided", u16::MAX)), + ))? + .queue(cursor::MoveToColumn(0))?; + } else { + drawn = true; + } + + let mut iter = choices.iter().enumerate().peekable(); + while let Some((i, choice)) = iter.next() { + // if active choice, flip foreground and background + // if active choice, wrap in [] + // if not, wrap in spaces, to preserve spacing and prevent redraws + if i == active_choice { + terminal.queue(PrintStyledContent(Stylize::reverse(format!("[{choice}]"))))?; + } else { + terminal.queue(Print(format!(" {choice} ")))?; + } + if iter.peek().is_some() { + if horizontal { + terminal.queue(Print(" "))?; + } else { + terminal.queue(cursor::MoveToNextLine(1))?; + } + } + } + terminal.flush()?; + + if let Event::Key(k) = read()? { + match k.code { + KeyCode::Char('c') if k.modifiers.contains(KeyModifiers::CONTROL) => { + return Err(Error::CtrlC); + } + KeyCode::Char(c) => { + for (i, choice) in choices.iter().enumerate() { + if choice.identifier().is_some_and(|id| id == c) { + active_choice = i; + } + } + } + KeyCode::Left | KeyCode::Up => { + active_choice = active_choice.saturating_sub(1); + } + KeyCode::Right | KeyCode::Down => match choices.len().saturating_sub(active_choice) { + 0 | 1 => {} + _ => { + active_choice += 1; + } + }, + KeyCode::Enter => { + return Ok(active_choice); + } + _ => {} + } + } + } + } + fn prompt_validated_wordlist( &mut self, prompt: &str, @@ -307,9 +407,7 @@ where prefix_length = line.len(); terminal.queue(Print(line))?; if lines.peek().is_some() { - terminal - .queue(cursor::MoveDown(1))? - .queue(cursor::MoveToColumn(0))?; + terminal.queue(cursor::MoveToNextLine(1))?; } } terminal.flush()?; @@ -468,9 +566,7 @@ where prefix_length = line.len(); terminal.queue(Print(line))?; if lines.peek().is_some() { - terminal - .queue(cursor::MoveDown(1))? - .queue(cursor::MoveToColumn(0))?; + terminal.queue(cursor::MoveToNextLine(1))?; } } terminal.flush()?; @@ -536,21 +632,17 @@ where let len = std::cmp::min(u16::MAX as usize, word.len()) as u16; written_chars += len + 1; if written_chars > cols { - terminal - .queue(cursor::MoveDown(1))? - .queue(cursor::MoveToColumn(0))?; + terminal.queue(cursor::MoveToNextLine(1))?; written_chars = len + 1; } terminal.queue(Print(word))?.queue(Print(" "))?; } - terminal - .queue(cursor::MoveDown(1))? - .queue(cursor::MoveToColumn(0))?; + terminal.queue(cursor::MoveToNextLine(1))?; } } Message::Data(data) => { let count = data.lines().count(); - // NOTE: GE to allow a MoveDown(1) + // NOTE: GE to allow a MoveToNextLine(1) if count >= rows as usize { let msg = format!( "{} {count} {} {rows} {}", @@ -558,14 +650,12 @@ where ); terminal .queue(Print(msg))? - .queue(cursor::MoveDown(1))? - .queue(cursor::MoveToColumn(0))?; + .queue(cursor::MoveToNextLine(1))?; } else { for line in data.lines() { terminal .queue(Print(line))? - .queue(cursor::MoveDown(1))? - .queue(cursor::MoveToColumn(0))?; + .queue(cursor::MoveToNextLine(1))?; } } } @@ -587,7 +677,6 @@ where _ => (), } } - terminal.queue(cursor::EnableBlinking)?.flush()?; Ok(()) } }