From 0d39dae09560030780af6d646969bbd65050ab47 Mon Sep 17 00:00:00 2001
From: Unbewohnte <unbewohnte@gmail.com>
Date: Sat, 1 Jan 2022 14:34:35 +0300
Subject: [PATCH] Safety additions. Do not follow links, track only regular
 files

---
 src/broom.cpp | 38 ++++++++++++++++++++++++--------------
 src/broom.hpp |  3 ++-
 src/main.cpp  | 14 +++++++++-----
 3 files changed, 35 insertions(+), 20 deletions(-)

diff --git a/src/broom.cpp b/src/broom.cpp
index 3ba3e39..8da0327 100644
--- a/src/broom.cpp
+++ b/src/broom.cpp
@@ -40,23 +40,34 @@ void Broom::print_statistics() {
     << std::endl;
 };
 
-// get all files from path recursively and track them
-void Broom::track(const std::filesystem::path dir) {
+// recursively track every file that lies in given path. Throws an invalid_argument
+// error in case path does not exist
+void Broom::track(const std::filesystem::path path) {
     auto t0 = std::chrono::high_resolution_clock::now();
 
-    std::filesystem::directory_options options =  (
-        std::filesystem::directory_options::follow_directory_symlink |
-        std::filesystem::directory_options::skip_permission_denied
-    );
+    // check if given path even exists
+    if (!std::filesystem::exists(path)) {
+        throw std::invalid_argument("\"" + path.string() + "\"" + " does not exist !");
+    };
 
-    for (std::filesystem::directory_entry dir_entry : std::filesystem::recursive_directory_iterator(dir, options)) {
-        if (dir_entry.is_directory()) {
-            continue;
-        };
+    if (std::filesystem::is_directory(path)) {
+        // it`s a directory. Track every regular file recursively
+        std::filesystem::directory_options options =  (
+            std::filesystem::directory_options::skip_permission_denied
+        );
 
-        m_tracked_filepaths.push_back(dir_entry);
+        for (std::filesystem::directory_entry dir_entry : std::filesystem::recursive_directory_iterator(path, options)) {
+            if (!dir_entry.is_regular_file()) {
+                continue;
+            };
+
+            m_tracked_filepaths.push_back(dir_entry);
+        };
+    } else if (std::filesystem::is_regular_file(path)) {
+        m_tracked_filepaths.push_back(path);
     };
 
+
     if (m_benchmarking) {
         auto tracking_time = std::chrono::high_resolution_clock::now();
 
@@ -78,6 +89,7 @@ uintmax_t Broom::untrack_unique_sizes() {
         // if yes --> increment occurences counter
         // if not --> add it to the map with a counter of 1
         uintmax_t filesize = std::filesystem::file_size(filepath);
+
         auto iterator = sizes_map.find(filesize);
         if (iterator == sizes_map.end()) {
             // there is no such size
@@ -161,9 +173,7 @@ void Broom::find_duplicates() {
         std::cout
         << "Untracking by size took "
         << std::chrono::duration_cast<std::chrono::milliseconds>(sizes_untrack_time - t0).count()
-        << " ms" << std::endl
-
-        << std::endl;
+        << " ms" << std::endl;
     } else {
         size_t startsize = m_tracked_filepaths.size();
         std::cout << "Tracking " << startsize << std::endl;
diff --git a/src/broom.hpp b/src/broom.hpp
index b723f4b..90d92a8 100644
--- a/src/broom.hpp
+++ b/src/broom.hpp
@@ -51,7 +51,8 @@ public:
     // Print current statistics
     void print_statistics();
 
-    // recursively track every file that lies in given path
+    // recursively track every file that lies in given path. Throws an invalid_argument
+    // error in case path does not exist
     void track(const std::filesystem::path path);
 
     // find all duplicates in the directory
diff --git a/src/main.cpp b/src/main.cpp
index cf99ae0..d39c846 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -18,6 +18,7 @@ along with broom.  If not, see <https://www.gnu.org/licenses/>.
 */
 
 #include <iostream>
+#include <stdexcept>
 #include <string.h>
 #include <vector>
 
@@ -82,13 +83,11 @@ int main(int argc, char* argv[]) {
         }
         else {
             // add path
-            if (std::filesystem::exists(argv[i])) {
-                tracked_path = argv[i];
-            };
+            tracked_path = argv[i];
         };
     };
 
-    // no path was specified at all or every path was nonexistent
+    // no path was specified at all
     if (tracked_path.string() == "") {
         print_help();
         return 1;
@@ -97,7 +96,12 @@ int main(int argc, char* argv[]) {
 
     Broom broom(options);
 
-    broom.track(tracked_path);
+    try {
+        broom.track(tracked_path);
+    } catch(const std::invalid_argument& e) {
+        std::cout << e.what() << std::endl;
+        return 1;
+    };
     broom.find_duplicates();
 
     return 0;