From a057cf1bd5586fe5f3605d661a5f56a9cdd7095e Mon Sep 17 00:00:00 2001
From: Alexandre CHAZAL <alexandre.chazal@epita.fr>
Date: Thu, 2 Sep 2021 20:08:23 +0200
Subject: [PATCH] fix(add): now working as intended

---
 src/add-sshkey.cc | 33 +++++++++++++++++++++++----------
 1 file changed, 23 insertions(+), 10 deletions(-)

diff --git a/src/add-sshkey.cc b/src/add-sshkey.cc
index 04f02e3..6548585 100644
--- a/src/add-sshkey.cc
+++ b/src/add-sshkey.cc
@@ -1,19 +1,22 @@
 #include <iostream>
 #include <fstream>
+#include <stdlib.h>
 
-int append_key(std::string &key, std::fstream &file)
+int append_key(std::string &key, std::fstream &f)
 {
     std::string tmp;
 
-    while (getline(file, tmp))
+    while (std::getline(f, tmp))
     {
         if (tmp == key)
             return 1;
     }
 
-    file << "\n" << key;
+    f.clear();
+    f.seekp(std::ios::end);
+    f << key << "\n";
 
-    if (file.bad())
+    if (f.bad())
         return 2;
 
     return 0;
@@ -30,7 +33,17 @@ int main(int argc, char *argv[])
 
     std::string key(argv[1]);
 
-    std::fstream file("~/.ssh/authorized_keys");
+    std::string authorized_keys = std::string(getenv("HOME"));
+    authorized_keys += "/.ssh/authorized_keys";
+
+    std::fstream file;
+    file.open(authorized_keys, std::ios::in | std::ios::out | std::ios::app);
+
+    if (!file)
+    {
+        std::cerr << "[!!] Could not open file, exiting." << std::endl;
+        return 2;
+    }
 
     int res = append_key(key, file);
     file.close();
@@ -39,15 +52,15 @@ int main(int argc, char *argv[])
     {
         std::cerr << "[!!] The key has already been added, exiting."
                   << std::endl;
+        return 3;
     }
-    else if (res == 2)
+    if (res == 2)
     {
         std::cerr << "[!!] Error while writing key, exiting." << std::endl;
+        return 4;
     }
-    else
-    {
-        std::cerr << "[++] Key has been added." << std::endl;
-    }
+
+    std::cerr << "[++] Key has been added." << std::endl;
 
     return 0;
 }