4 files changed, 34 insertions, 20 deletions

diff --git a/src/flags.rs b/src/flags.rs
index 91acf3a..9432b8f 100644
--- a/src/flags.rs
+++ b/src/flags.rs
@@ -5,7 +5,7 @@ pub struct Flags {
   pub arg_count: usize
 }
 
-pub fn parse(args: &[String]) -> Flags {
+pub fn parse(args: &[String]) -> Option<Flags> {
   let mut flags = Flags { 
     help: false, 
     version: false,
@@ -17,15 +17,21 @@ pub fn parse(args: &[String]) -> Flags {
     flags.arg_count += 1;
     if arg.starts_with("--") {
       let flag = &arg[2..];
-      check_flag(&flag, &mut flags);
+      if !set_flag(&flag, &mut flags) {
+        eprintln!("Invalid argument: {}", arg);
+        return None
+      }
     } else {
       let flag = &arg[1..];
       for char in flag.chars() {
-        check_flag(&char.to_string(), &mut flags);
+        if !set_flag(&char.to_string(), &mut flags) {
+            eprintln!("Invalid argument: {}", arg);
+            return None
+        }
       }
     }
   }
-  flags
+  Some(flags)
 }
 
 fn is_arg(arg: &str) -> bool {
@@ -36,16 +42,18 @@ const HELP_FLAG:              &str = "help h";
 const VERSION_FLAG:           &str = "version v";
 const DONT_PERSIST:           &str = "d";
 
-fn check_flag(arg: &str, flags: &mut Flags) {
+fn set_flag(arg: &str, flags: &mut Flags) -> bool {
   if has_flag_set(&arg, HELP_FLAG) {
-    flags.help = true
-  }
-  if has_flag_set(&arg, VERSION_FLAG) {
-    flags.version = true
-  }
-  if has_flag_set(&arg, DONT_PERSIST) {
-    flags.dont_persist = true
+    flags.help = true;
+    return true
+  } else if has_flag_set(&arg, VERSION_FLAG) {
+    flags.version = true;
+    return true
+  } else if has_flag_set(&arg, DONT_PERSIST) {
+    flags.dont_persist = true;
+    return true
   }
+  false
 }
 
 fn has_flag_set(arg: &str, check: &str) -> bool {
@@ -55,4 +63,4 @@ fn has_flag_set(arg: &str, check: &str) -> bool {
     }
   }
   return false
-}
\ No newline at end of file
+}
diff --git a/src/help.rs b/src/help.rs
index 422e35d..e6f4a72 100644
--- a/src/help.rs
+++ b/src/help.rs
@@ -2,10 +2,9 @@ pub fn help() {
   let help = 
 "Usage:
     crab [-d] command [args]
-
 Options:
     -v --version    Get the current version of the package
     -h --help       Generates the crab help message
     -d              If your user is set to persist, dont save persistance";
   println!("{}", help);
-}
\ No newline at end of file
+}
diff --git a/src/main.rs b/src/main.rs
index 5fab1d9..7c3829f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,7 +5,7 @@ use nix::unistd;
 
 extern crate time;
 
-// const ERROR_COMMAND:           u8 = 1;
+const ERROR_ARGS:           u8 = 1;
 const ERROR_CONFIG:         u8 = 2;
 const ERROR_NO_USER:        u8 = 3;
 const ERROR_NOT_AUTHORIZED: u8 = 4;
@@ -18,7 +18,13 @@ mod help;
 
 fn main() -> ExitCode {
     let args: Vec<String> = env::args().collect();
-    let flags = flags::parse(&args[1..]);
+    let flags = match flags::parse(&args[1..]) {
+        Some(data) => data,
+        None => {
+            help::help();
+            return ExitCode::from(ERROR_ARGS);
+        }
+    };
     if flags.version {
         println!("crab version 0.0.3");
         return ExitCode::SUCCESS;
@@ -132,4 +138,4 @@ fn config(path: &str) -> Option<Config> {
         users.push((user, persist));
     }
     Some(Config{users})
-}
\ No newline at end of file
+}
diff --git a/src/persist.rs b/src/persist.rs
index 7d47232..febd035 100644
--- a/src/persist.rs
+++ b/src/persist.rs
@@ -83,12 +83,13 @@ fn get_terminal_config() -> Option<Value> {
 }
 
 fn write_terminal_config(id: &i32, data: &str) -> Result<(), Box<dyn std::error::Error>> {
-  std::fs::write(path(&id), data)?;
+  std::fs::write(path(&id), "")?;
   unistd::chown(std::path::Path::new(&path(&id)), Some(unistd::Uid::from(0)), Some(unistd::Gid::from(0)))?;
   let metadata = std::fs::metadata(path(&id))?;
   let mut perms = metadata.permissions();
   perms.set_mode(0o660);
   fs::set_permissions(path(&id), perms)?;
+  std::fs::write(path(&id), data)?;
   Ok(())
 }
 
@@ -98,4 +99,4 @@ fn now() -> u64 {
 
 fn path(id: &i32) -> String {
   return format!("/tmp/crab-{}", id);
-}
\ No newline at end of file
+}