dotfiles/scripts/health/health.c

#define DEBUG

#include <stdlib.h>
#include <stdio.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fts.h>
#include <string.h>
#include <errno.h>
#include <math.h>
#include <stdarg.h>

#include "health.h"
#include "utils.h"

//ustring ok_str = NEW_USTRING("OK");
//ustring warn_str = NEW_USTRING("WARN");
//ustring arrow_str = NEW_USTRING("├─");
//ustring arrow_last_str = NEW_USTRING("└─");
ustring ok_str = {.str = "OK", .bytes = strlen("OK") };
ustring warn_str = {.str = "WARN ☠", .bytes = strlen("WARN ☠")};
ustring arrow_str = {.str = "├─", .bytes = strlen("├─")};
ustring arrow_last_str = {.str ="└─", .bytes = strlen("└─")};

// fts_path doesn't append trailing slash
struct dir_status top_dir_var[] = {
    WATCH_DIR("/var/cache", "1.5G"),
    WATCH_DIR("/var/log", "0.5G")
};

struct dir_status top_dir_home[] = {
    WATCH_DIR("/home/kuba/.cache", "3G"),
    WATCH_DIR("/home/kuba/Downloads", "1G")
};

struct dir_status top_dir_usr[] = {
    WATCH_DIR("/usr/local", "5G"),
    WATCH_DIR("/usr/lib", "5G"),
    WATCH_DIR("/usr/share", "5G"),
    WATCH_DIR("/usr/bin", "1G")
};

struct dir_status watched_dirs[] = {
    WATCH_DIR_SUBS("/var", "2.5G", top_dir_var),
    WATCH_DIR_SUBS("/home/kuba", "20G", top_dir_home),
    WATCH_DIR_SUBS("/usr", "15G", top_dir_usr),
    WATCH_DIR("/boot", "80M"),
    WATCH_DIR("/etc", "20M"),
    WATCH_DIR("/root", "10M"),
    WATCH_DIR("/run", "10M"),
    WATCH_DIR("/srv", "1M"),
    WATCH_DIR("/opt", "1G")
};

int main(int argc, char* const argv[])
{
    INIT_USTRING(ok_str);
    INIT_USTRING(warn_str);
    INIT_USTRING(arrow_str);
    INIT_USTRING(arrow_last_str);
    // Traverse top directories and do some nice formatting
    char buf[80];
    for (int i = 0; i < sizeof(watched_dirs)/sizeof(struct dir_status); i++){
        struct dir_status *dir = watched_dirs + i;
        disk_usage(dir);

        format_directory_entry(buf, dir, 0, 0);
    }
    exit(0);

    if (argc<2)
    {
        printf("Usage: %s <path-spec>\n", argv[0]);
        exit(255);
    }

    return 0;
}


//Maybe just remove this, the effect is not noticable
off_t traverse_nochildren(FTS *file_system, const short dir_level)
{
    FTSENT *node     = NULL;
    struct stat *st  = NULL;
    off_t tot_size = 0;

    while( (node = fts_read(file_system)) != NULL)
    {
        // Add to size
        st = node->fts_statp;
        if (st) tot_size += (off_t) 512*st->st_blocks;
        else debug("st null pointer, skipping\n");

        if (node->fts_info == FTS_DP)
        {
            indent_debug(node->fts_level);
            debug("Exiting %s\n", node->fts_path);
            // Directory being visited in post-order. Only action is to
            // return if back at the starting level
            if (node->fts_level == dir_level) return tot_size;
        }
    }
    return 0;// This is bad
}

void traverse_children(FTS *file_system, struct dir_status *dir, const short dir_level)
{
    FTSENT *node     = NULL;
    struct stat *st  = NULL;
    off_t tot_size = 0;

    top_loop:while( (node = fts_read(file_system)) != NULL)
    {
        //indent_debug(node->fts_level);
        //debug("%s\n", node->fts_name);
        if (node->fts_info == FTS_D)
        {
            if(strcmp(node->fts_path, dir->path) == 0) {
                continue; // Do nothing
            }
            indent_debug(node->fts_level);
            debug("Entering %s\n", node->fts_path);
            // Directory being visited in pre-order. Do not count size except
            // after returning from a recursion step. Determine if the current
            // node is one of the children.
            for (int i = 0; i < dir->n_children; i++) {
                struct dir_status *sub_dir = dir->children+i;
                if(strcmp(node->fts_path, sub_dir->path) == 0){
                    // Entering one of the child-directories
                    indent_debug(node->fts_level);
                    debug("Entering special %s\n", node->fts_path);
                    traverse_children(file_system, sub_dir, node->fts_level);
                    tot_size += sub_dir->size;
                    goto top_loop;
                }
            }
            // Child directory not in list, call traverse_nochildren instead
            tot_size += traverse_nochildren(file_system, node->fts_level);
            continue; // Does the same as goto above
        }

        // Add to size
        st = node->fts_statp;
        if (st) tot_size += (off_t) 512*st->st_blocks;
        else debug("st null pointer, skipping\n");

        if (node->fts_info == FTS_DP)
        {
            indent_debug(node->fts_level);
            debug("Exiting %s\n", node->fts_path);
            // Directory being visited in post-order. Only action is to save
            // size and return if back at the starting level
            if (node->fts_level == dir_level){
                dir->size = tot_size;
                return;
            }
        }
    }
}

int disk_usage(struct dir_status *top_dir)
{
    char *paths[] = { top_dir->path, NULL};
    FTS* file_system = NULL;

    file_system = fts_open(paths, FTS_PHYSICAL, NULL);

    if (file_system == NULL) return 1;

    traverse_children(file_system, top_dir, 0);

    fts_close(file_system);
    return 0;
}

const int max_length = 60;
const int max_path_length = 40;
void format_directory_entry(char *buf, struct dir_status *dir, int level, int last_subdir)
{
    int n = 0; //Keeps track of where in buffer we are
    int delta_n = 0; // Keeps track of how much shorter the string appears
                     // than the number of bytes

    // First fill with apropriate amount of spaces and └ or ├ charcters
    if (level > 0) {
        ustring *arr = last_subdir?&arrow_last_str:&arrow_str;

        n = 4*level - arr->chars;
        delta_n += arr->chars - arr->bytes;
        memset(buf, ' ', n);
        strncpy(buf+n, arr->str, arr->bytes);
        n += arr->bytes;
    }

    // Path and size
    n += snprintf(buf + n, max_path_length, "%s -- ", dir->path);
    n += pretty_bytes(buf + n, dir->size);
    memset(buf + n, ' ', max_length - n); // max_length - n for simplicity

    // Right-justify WARN or OK status
    ustring *ok = directory_ok(dir->size, dir->warn_at)?&ok_str:&warn_str;
    delta_n += ok->chars - ok->bytes;
    snprintf(buf + max_length - ok->bytes - 2 - delta_n, ok->bytes+2
            , "%s\n", ok->str);// 2 leaves space for \n and null
    printf("%s", buf);

    // Loop through children
    for (int i = 0; i < dir->n_children; i++){
        int last = (i == dir->n_children - 1);
        format_directory_entry(buf, dir->children + i, (level + 1), last);
    }

    if(level == 0) printf("\n");
}
int directory_ok(off_t size, char *warn_at){
    off_t warn = pretty_bytes_to_num(warn_at);
    return size < warn;
}