From 39ba4c9087eda99c44ee4c7163d1ae1a4435bff4 Mon Sep 17 00:00:00 2001
From: apio <blobs.trading@gmail.com>
Date: Sat, 2 Sep 2023 20:01:10 +0200
Subject: [PATCH] ls: Add colors to output

---
 apps/ls.cpp | 66 +++++++++++++++++++++++++++++++++++++++++++++++------
 1 file changed, 59 insertions(+), 7 deletions(-)

diff --git a/apps/ls.cpp b/apps/ls.cpp
index b5d4fe58..ed0141ba 100644
--- a/apps/ls.cpp
+++ b/apps/ls.cpp
@@ -8,6 +8,7 @@
 #include <os/FileSystem.h>
 #include <os/Mode.h>
 #include <pwd.h>
+#include <unistd.h>
 
 void find_user_and_group(struct stat& st, StringView& owner, StringView& group)
 {
@@ -47,18 +48,46 @@ int sort_reverse(const os::Directory::Entry* a, const os::Directory::Entry* b)
     return 0;
 }
 
-static Result<String> entry_join(const Vector<os::Directory::Entry>& vec, StringView delim)
+#define RESET_COLORS "\x1b[m"
+#define SYMLINK_COLOR "\x1b[36m"
+#define FILE_COLOR "\x1b[1;32m"
+#define DIR_COLOR "\x1b[1;34m"
+#define SOCKET_COLOR "\x1b[33m"
+#define SPECIAL_COLOR "\x1b[35m"
+#define STICKY_COLOR "\x1b[30;1;42m"
+#define SETUID_COLOR "\x1b[30;1;41m"
+#define EXEC_COLOR "\x1b[1;31m"
+
+static const char* file_type_color(const os::Directory::Entry& entry)
+{
+    if (entry.mode & S_ISVTX) return STICKY_COLOR;
+    if (entry.mode & S_ISUID || entry.mode & S_ISGID) return SETUID_COLOR;
+
+    switch (entry.mode & S_IFMT)
+    {
+    case S_IFREG: return entry.mode & S_IXUSR ? EXEC_COLOR : FILE_COLOR;
+    case S_IFDIR: return DIR_COLOR;
+    case S_IFLNK: return SYMLINK_COLOR;
+    case S_IFSOCK: return SOCKET_COLOR;
+    default: return SPECIAL_COLOR;
+    }
+}
+
+static Result<String> entry_join(const Vector<os::Directory::Entry>& vec, StringView delim, bool colors)
 {
     if (vec.size() == 0) return String {};
-    if (vec.size() == 1) return vec[0].name.clone();
 
     StringBuilder sb;
+    if (colors) TRY(sb.add(StringView { file_type_color(vec[0]) }));
     TRY(sb.add(vec[0].name));
+    if (colors) TRY(sb.add(StringView { RESET_COLORS }));
 
     for (usize i = 1; i < vec.size(); i++)
     {
         TRY(sb.add(delim));
+        if (colors) TRY(sb.add(StringView { file_type_color(vec[i]) }));
         TRY(sb.add(vec[i].name));
+        if (colors) TRY(sb.add(StringView { RESET_COLORS }));
     }
 
     return sb.string();
@@ -75,6 +104,7 @@ Result<int> luna_main(int argc, char** argv)
     bool follow_symlink_args { false };
     bool one_per_line { false };
     bool list_directories { false };
+    bool no_colors { false };
 
     StringView sort_type { "name" };
 
@@ -98,6 +128,8 @@ Result<int> luna_main(int argc, char** argv)
     parser.add_switch_argument(list_directories, 'd', "directory"_sv, "list directories instead of their contents"_sv);
     parser.add_value_argument(sort_type, ' ', "sort"_sv, "sort by name, size or time"_sv);
     parser.add_switch_argument(reverse_sort, 'r', "reverse"_sv, "reverse order while sorting"_sv);
+    parser.add_switch_argument(no_colors, ' ', "no-colors"_sv,
+                               "disable coloring of output (defaults to true when not in a TTY)"_sv);
     parser.parse(argc, argv);
 
     Vector<os::Directory::Entry> files;
@@ -154,11 +186,13 @@ Result<int> luna_main(int argc, char** argv)
 
     if (!long_list)
     {
-        auto list = TRY(entry_join(files, one_per_line ? "\n"_sv : " "_sv));
+        auto list = TRY(entry_join(files, one_per_line ? "\n"_sv : " "_sv, !no_colors && isatty(STDIN_FILENO)));
         if (!list.is_empty()) os::println("%s", list.chars());
     }
     else
     {
+        bool colors = !no_colors && isatty(STDIN_FILENO);
+
         for (const auto& file : files)
         {
             struct stat st;
@@ -176,14 +210,32 @@ Result<int> luna_main(int argc, char** argv)
 
             if (!human_readable && !si)
             {
-                os::println("%s %u %4s %4s %10lu %s%s%s", formatted_mode, st.st_nlink, owner.chars(), group.chars(),
-                            st.st_size, file.name.chars(), link.is_empty() ? "" : " -> ", link.chars());
+                if (colors)
+                {
+                    os::println("%s %u %4s %4s %10lu %s%s" RESET_COLORS "%s" SYMLINK_COLOR "%s" RESET_COLORS,
+                                formatted_mode, st.st_nlink, owner.chars(), group.chars(), st.st_size,
+                                file_type_color(file), file.name.chars(), link.is_empty() ? "" : " -> ", link.chars());
+                }
+                else
+                {
+                    os::println("%s %u %4s %4s %10lu %s%s%s", formatted_mode, st.st_nlink, owner.chars(), group.chars(),
+                                st.st_size, file.name.chars(), link.is_empty() ? "" : " -> ", link.chars());
+                }
             }
             else
             {
                 auto size = TRY(to_dynamic_unit(st.st_size, 10, false, si ? Unit::SI : Unit::Binary, false));
-                os::println("%s %u %4s %4s %6s %s%s%s", formatted_mode, st.st_nlink, owner.chars(), group.chars(),
-                            size.chars(), file.name.chars(), link.is_empty() ? "" : " -> ", link.chars());
+                if (colors)
+                {
+                    os::println("%s %u %4s %4s %6s %s%s" RESET_COLORS "%s" SYMLINK_COLOR "%s" RESET_COLORS,
+                                formatted_mode, st.st_nlink, owner.chars(), group.chars(), size.chars(),
+                                file_type_color(file), file.name.chars(), link.is_empty() ? "" : " -> ", link.chars());
+                }
+                else
+                {
+                    os::println("%s %u %4s %4s %6s %s%s%s", formatted_mode, st.st_nlink, owner.chars(), group.chars(),
+                                size.chars(), file.name.chars(), link.is_empty() ? "" : " -> ", link.chars());
+                }
             }
         }
     }