1 files changed, 29 insertions, 25 deletions
diff --git a/src/prompt.rs b/src/prompt.rs
index fd9c872..e6600b7 100644
--- a/src/prompt.rs
+++ b/src/prompt.rs
@@ -1,22 +1,12 @@
 use std::io::Write;
 
 pub struct Prompt {
-    home: std::path::PathBuf,
-    pub pwd: std::path::PathBuf,
-    username: std::string::String,
-    hostname: std::string::String,
+    ps1: std::string::String,
+    pub home: std::path::PathBuf,
 }
 
 impl Prompt {
     pub fn new() -> Prompt {
-        let home = match std::env::home_dir() {
-            Some(p) => p,
-            None => std::path::Path::new("/").to_path_buf(),
-        };
-        let pwd = match std::env::current_dir() {
-            Ok(p) => p,
-            Err(e) => panic!(e),
-        };
         let username = match std::process::Command::new("whoami").output() {
             Ok(s) => {
                 let output = std::string::String::from_utf8(s.stdout).unwrap();
@@ -31,24 +21,38 @@ impl Prompt {
             }
             Err(e) => panic!(e),
         };
-        Prompt {
-            home,
-            pwd,
-            username,
-            hostname,
-        }
+        let ready = if username == "root" {
+            "#"
+        } else {
+            "$"
+        }.to_owned();
+
+        let ps1 = String::from("  {USER}@{HOST} [{PWD}]\n\\$ ")
+            .replace("{USER}", &username)
+            .replace("{HOST}", &hostname)
+            .replace("\\$", &ready);
+        let home = match std::env::home_dir() {
+            Some(p) => p,
+            None => std::path::Path::new("/").to_path_buf(),
+        };
+
+        Prompt { ps1, home }
     }
 
-    pub fn print(&self) {
-        let path = if self.pwd == self.home {
-            self.pwd.display()
-        } else if self.pwd.starts_with(&self.home) {
-            self.pwd.strip_prefix(&self.home).unwrap().display()
+    pub fn print(&self, pwd: &std::path::PathBuf) {
+        if cfg!(debug_assertions) {
+            println!("pwd={}", pwd.display());
+        }
+
+        let path = if pwd.starts_with(&self.home) {
+            "~".to_owned() + &pwd.strip_prefix(&self.home).unwrap().display().to_string()
         } else {
-            self.pwd.display()
+            pwd.display().to_string()
         };
 
-        print!("  {USER}@{HOST} [{PWD}]\n$ ", USER=self.username, HOST=self.hostname, PWD=path);
+        let ps1 = self.ps1.replace("{PWD}", &path);
+
+        print!("{}", ps1);
         match std::io::stdout().flush() {
             Ok(_) => {}
             Err(_) => println!(""),