From b6e838bab0ef49c84caa379a051cf6129becf979 Mon Sep 17 00:00:00 2001
From: Roey Darwish Dror <roey.ghost@gmail.com>
Date: Sun, 16 Jun 2019 09:09:05 +0300
Subject: [PATCH] Add -k flag

---
 src/config.rs        | 11 ++++++++++-
 src/main.rs          |  4 ++--
 src/steps/os/unix.rs |  3 ++-
 src/terminal.rs      |  2 +-
 4 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 0bf5df6d..97fcfe65 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -5,7 +5,7 @@ use lazy_static::lazy_static;
 use serde::Deserialize;
 use shellexpand;
 use std::collections::{BTreeMap, HashMap};
-use std::fs;
+use std::{env, fs};
 use structopt::StructOpt;
 use toml;
 
@@ -125,6 +125,10 @@ pub struct CommandLineArgs {
     /// Output logs
     #[structopt(short = "v", long = "verbose")]
     verbose: bool,
+
+    /// Prompt or a key before exiting
+    #[structopt(short = "k", long = "keep")]
+    keep_at_end: bool,
 }
 
 /// Represents the application configuration
@@ -207,4 +211,9 @@ impl Config {
     pub fn remote_topgrades(&self) -> &Option<Vec<String>> {
         &self.config_file.remote_topgrades
     }
+
+    /// Prompt for a key before exiting
+    pub fn keep_at_end(&self) -> bool {
+        self.opt.keep_at_end || env::var("TOPGRADE_KEEP_END").is_ok()
+    }
 }
diff --git a/src/main.rs b/src/main.rs
index 9b4e26d7..cbc8b5aa 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -498,8 +498,8 @@ fn run() -> Result<(), Error> {
         freebsd::audit_packages(&sudo).ok();
     }
 
-    if env::var("TOPGRADE_KEEP_END").is_ok() {
-        print_info("\nPress R to reboot, S for shell or any other key to continue");
+    if config.keep_at_end() {
+        print_info("\n(R)eboot\n(S)hell\n\nPress any other key to continue");
         match get_char() {
             's' | 'S' => {
                 run_shell();
diff --git a/src/steps/os/unix.rs b/src/steps/os/unix.rs
index 3a771ea8..162c8de8 100644
--- a/src/steps/os/unix.rs
+++ b/src/steps/os/unix.rs
@@ -131,5 +131,6 @@ pub fn run_sdkman(base_dirs: &BaseDirs, cleanup: bool, run_type: RunType) -> Res
 }
 
 pub fn reboot() {
-    Command::new("sudo").arg("reboot").spawn().ok();
+    print!("Rebooting...");
+    Command::new("sudo").arg("reboot").spawn().unwrap().wait().unwrap();
 }
diff --git a/src/terminal.rs b/src/terminal.rs
index 74f96205..7fe4a502 100644
--- a/src/terminal.rs
+++ b/src/terminal.rs
@@ -92,7 +92,7 @@ impl Terminal {
     fn print_info<P: AsRef<str>>(&mut self, message: P) {
         let message = message.as_ref();
         self.term
-            .write_fmt(format_args!("{}\n", style(message).yellow().bold()))
+            .write_fmt(format_args!("{}\n", style(message).blue().bold()))
             .ok();
     }