From ca86b4464031d9a91a289917b8382bf784c8ad9e Mon Sep 17 00:00:00 2001
From: Carson McManus <carson.mcmanus1@gmail.com>
Date: Sat, 14 Aug 2021 11:10:21 -0400
Subject: [PATCH] add prompt_char helper function

---
 src/main.rs | 18 ++++++------
 src/tui.rs  | 82 ++++++++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 89 insertions(+), 11 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 7e58f5a..7a58edb 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -125,12 +125,11 @@ fn main() {
 	let mut manifest: accountmanager::Manifest;
 	if !path.exists() {
 		error!("Did not find manifest in {}", mafiles_dir);
-		print!(
-			"Would you like to create a manifest in {} ? [Yn] ",
-			mafiles_dir
-		);
-		match tui::prompt().to_lowercase().as_str() {
-			"n" => {
+		match tui::prompt_char(
+			format!("Would you like to create a manifest in {} ?", mafiles_dir).as_str(),
+			"Yn",
+		) {
+			'n' => {
 				info!("Aborting!");
 				return;
 			}
@@ -352,11 +351,10 @@ fn main() {
 				.join(", ")
 		);
 
-		print!("Do you want to continue? [yN] ");
-		match tui::prompt().as_str() {
-			"y" => {}
+		match tui::prompt_char("Do you want to continue?", "yN") {
+			'y' => {}
 			_ => {
-				println!("Aborting!");
+				info!("Aborting!");
 				return;
 			}
 		}
diff --git a/src/tui.rs b/src/tui.rs
index 4094969..bdabc11 100644
--- a/src/tui.rs
+++ b/src/tui.rs
@@ -1,7 +1,7 @@
 use log::*;
 use regex::Regex;
 use std::collections::HashSet;
-use std::io::Write;
+use std::io::{Read, Write};
 use steamguard::Confirmation;
 use termion::{
 	event::{Event, Key},
@@ -61,6 +61,46 @@ pub fn prompt_captcha_text(captcha_gid: &String) -> String {
 	return captcha_text;
 }
 
+/// Prompt the user for a single character response. Useful for asking yes or no questions.
+///
+/// `chars` should be all lowercase characters, with at most 1 uppercase character. The uppercase character is the default answer if no answer is provided.
+pub fn prompt_char(text: &str, chars: &str) -> char {
+	return prompt_char_impl(&mut std::io::stdin(), text, chars);
+}
+
+fn prompt_char_impl(input: &mut impl Read, text: &str, chars: &str) -> char {
+	let uppers = chars.replace(char::is_lowercase, "");
+	if uppers.len() > 1 {
+		panic!("Invalid chars for prompt_char. Maximum 1 uppercase letter is allowed.");
+	}
+	let default_answer: Option<char> = if uppers.len() == 1 {
+		Some(uppers.chars().collect::<Vec<char>>()[0].to_ascii_lowercase())
+	} else {
+		None
+	};
+
+	loop {
+		print!("{} [{}] ", text, chars);
+		let answer = input
+			.read_line()
+			.expect("Unable to read input")
+			.unwrap()
+			.to_ascii_lowercase();
+		if answer.len() == 0 {
+			if let Some(a) = default_answer {
+				return a;
+			}
+		} else if answer.len() > 1 {
+			continue;
+		}
+
+		let answer_char = answer.chars().collect::<Vec<char>>()[0];
+		if chars.contains(answer_char) {
+			return answer_char;
+		}
+	}
+}
+
 /// Returns a tuple of (accepted, denied). Ignored confirmations are not included.
 pub fn prompt_confirmation_menu(
 	confirmations: Vec<Confirmation>,
@@ -184,3 +224,43 @@ pub fn pause() {
 	stdout.flush().unwrap();
 	std::io::stdin().events().next();
 }
+
+#[cfg(test)]
+mod prompt_char_tests {
+	use super::*;
+
+	#[test]
+	fn test_gives_answer() {
+		let inputs = ['y', '\n'].iter().collect::<String>();
+		let answer = prompt_char_impl(&mut inputs.as_bytes(), "ligma balls", "yn");
+		assert_eq!(answer, 'y');
+	}
+
+	#[test]
+	fn test_gives_default() {
+		let inputs = ['\n'].iter().collect::<String>();
+		let answer = prompt_char_impl(&mut inputs.as_bytes(), "ligma balls", "Yn");
+		assert_eq!(answer, 'y');
+	}
+
+	#[test]
+	fn test_should_not_give_default() {
+		let inputs = ['n', '\n'].iter().collect::<String>();
+		let answer = prompt_char_impl(&mut inputs.as_bytes(), "ligma balls", "Yn");
+		assert_eq!(answer, 'n');
+	}
+
+	#[test]
+	fn test_should_not_give_invalid() {
+		let inputs = ['g', '\n', 'n', '\n'].iter().collect::<String>();
+		let answer = prompt_char_impl(&mut inputs.as_bytes(), "ligma balls", "yn");
+		assert_eq!(answer, 'n');
+	}
+
+	#[test]
+	fn test_should_not_give_multichar() {
+		let inputs = ['y', 'y', '\n', 'n', '\n'].iter().collect::<String>();
+		let answer = prompt_char_impl(&mut inputs.as_bytes(), "ligma balls", "yn");
+		assert_eq!(answer, 'n');
+	}
+}