/*
 * ftop - show progress of open files and file systems
 * Copyright (C) 2009  Jason Todd <jtodd1@earthlink.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/* snapshot.c - code to grab a "snapshot" of all currently open files */

#define _FILE_OFFSET_BITS 64
#include <sys/time.h>       /* gettimeofday */
#include <time.h>           /* gettimeofday */
#include <string.h>         /* memset, strerror, strncmp, memcpy, strrchr */
#include <glob.h>           /* glob, GLOB_NOSORT, globfree */
#include <stdlib.h>         /* free, malloc, strtoul, strtoull */
#include <stdio.h>          /* snprintf */
#include <errno.h>          /* errno, ENOENT */
#include <unistd.h>         /* readlink, getpid, stat, read, close */
#include <sys/types.h>      /* getpid, stat, open */
#include <sys/stat.h>       /* stat, S_IFIFO, S_IFSOCK, S_ISREG, S_ISDIR, */
                            /* S_ISCHR, S_ISBLK, S_ISFIFO, S_ISSOCK, open */
#include <fcntl.h>          /* open, O_RDONLY */

#include "ftop.h"
#include "snapshot.h"
#include "log.h"

static int snapshot_read_file(snapshot *s, char *file, size_t *bytes);
static int snapshot_buf_avail(snapshot *s, size_t buf_used, size_t bytes,
                              size_t *buf_size, size_t initial, char *name,
                              void *vbuf, snapshot_fixup_function f);
static int fix_file_paths(snapshot *s, char *new_buf);
static int fix_proc_cmds(snapshot *s, char *new_buf);
static int fix_proc_file(snapshot *s, char *new_buf);

int get_snapshot(snapshot *s)
{
    int ret;
    size_t i, j, read_bytes;
    ssize_t readlink_len;
    glob_t proc_glob, fd_glob;
    char *num_end, *num_str;
    char tmp_path[32];
    struct stat st;
    process_info *proc;
    file_info *file;

    gettimeofday(&(s->tv), NULL);
    s->num_processes = 0;
    s->num_unreadable_processes = 0;
    s->num_files = 0;
    s->num_reg = 0;
    s->num_dir = 0;
    s->num_chr = 0;
    s->num_blk = 0;
    s->num_pipe = 0;
    s->num_sock = 0;
    s->num_misc = 0;
    s->paths_buf_used = 0;
    s->log_used = 0;
    s->cmds_buf_used = 0;

    memset(&proc_glob, 0, sizeof(proc_glob));
    if (((ret = glob("/proc/*", GLOB_NOSORT, NULL, &proc_glob)) != 0) ||
        (proc_glob.gl_pathc == 0))
    {
        snapshot_log(s, LOG_ERR,
                     "could not list files in /proc/ (returned %d)", ret);

        return 0;
    }

    if (proc_glob.gl_pathc > s->alloc_processes)
    {
        snapshot_log(s, LOG_DEBUG,
                     "growing process info buffer from %u to %u processes",
                     s->alloc_processes, proc_glob.gl_pathc);

        if (s->proc != NULL)
            free(s->proc);

        s->proc = malloc(proc_glob.gl_pathc * sizeof(*(s->proc)));
        if (s->proc == NULL)
        {
            snapshot_log(s, LOG_ERR, "could not allocate process info buffer");
            return 0;
        }

        s->alloc_processes = proc_glob.gl_pathc;
    }

    if (s->readlink_buf == NULL)
    {
        s->readlink_buf_size = INITIAL_READLINK;
        if ((s->readlink_buf = malloc(s->readlink_buf_size)) == NULL)
        {
            snapshot_log(s, LOG_ERR, "could not allocate readlink buffer");
            return 0;
        }
    }

    for (i = 0; i < proc_glob.gl_pathc; i++)
    {
        proc = &(s->proc[s->num_processes]);
        memset(proc, 0, sizeof(*proc));

        num_str = proc_glob.gl_pathv[i] + 6;
        proc->pid = strtoul(num_str, &num_end, 0);
        if ((num_end == num_str) || (*num_end != '\0'))
        {
            snapshot_log(s, LOG_DEBUG, "skipping %s (not a PID)",
                         proc_glob.gl_pathv[i]);

            continue;
        }

        snprintf(tmp_path, sizeof(tmp_path), "/proc/%u/fd/*", proc->pid);
        memset(&fd_glob, 0, sizeof(fd_glob));
        errno = 0;
        if (glob(tmp_path, GLOB_NOSORT, NULL, &fd_glob) != 0)
        {
            if (errno != 0)
            {
                s->num_unreadable_processes++;
                snapshot_log(s, LOG_WARN,
                             "could not list files in /proc/%u/fd/ (%s)",
                             proc->pid, strerror(errno));
            }

            continue;
        }

        for (j = 0; j < fd_glob.gl_pathc; j++)
        {
            if (!snapshot_buf_avail(s, s->num_files * sizeof(*(s->file)),
                                    sizeof(*(s->file)), &(s->files_buf_size),
                                    INITIAL_FILES, "file info",
                                    &(s->file), fix_proc_file))
            {
                return 0;
            }

            file = &(s->file[s->num_files]);
            memset(file, 0, sizeof(*file));
            readlink_len = readlink(fd_glob.gl_pathv[j], s->readlink_buf,
                                    s->readlink_buf_size);

            while (readlink_len == s->readlink_buf_size)
            {
                snapshot_log(s, LOG_DEBUG,
                             "growing readlink buffer from %u to %u bytes",
                             s->readlink_buf_size, readlink_len << 1);

                free(s->readlink_buf);
                s->readlink_buf_size = readlink_len << 1;
                if ((s->readlink_buf = malloc(s->readlink_buf_size)) == NULL)
                {
                    snapshot_log(s, LOG_ERR,
                                 "could not allocate readlink buffer");

                    return 0;
                }

                readlink_len = readlink(fd_glob.gl_pathv[j], s->readlink_buf,
                                        s->readlink_buf_size);
            }

            if (readlink_len < 0)
            {
                if ((errno == ENOENT) && (proc->pid == getpid()))
                {
                    /*
                     * Non-existent glob result /proc/#/fd/3 (when # is "me")
                     * seems to be residual from the point in time when glob
                     * was doing its thing... specifically, it probably used to
                     * point to /proc/#/fd/ when glob had that directory open.
                     */
                    snapshot_log(s, LOG_INFO,
                                 "error reading link %s - just me, globbing myself",
                                 fd_glob.gl_pathv[j]);
                }
                else
                {
                    /* for all other cases (PID not "me"), it's an error */
                    snapshot_log(s, LOG_ERR, "error (%d) reading link %s",
                                 strerror(errno), fd_glob.gl_pathv[j]);
                }

                continue;
            }

            s->readlink_buf[readlink_len++] = '\0';
            if ((s->readlink_buf[0] != '/') ||
                (stat(s->readlink_buf, &(file->st)) < 0))
            {
                memset(&(file->st), 0, sizeof(file->st));
                if (s->readlink_buf[0] == '/')
                {
                    snapshot_log(s, LOG_INFO,
                                 "no stat info for '%s' in PID %u (%s)",
                                 s->readlink_buf, proc->pid, strerror(errno));

                    file->flags |= FLAG_MISC;
                }
                else
                {
                    if (strncmp(s->readlink_buf, "pipe:", 5) == 0)
                    {
                        file->st.st_mode = S_IFIFO;
                        s->num_pipe++;
                    }
                    else if (strncmp(s->readlink_buf, "socket:", 7) == 0)
                    {
                        file->st.st_mode = S_IFSOCK;
                        s->num_sock++;
                    }
                    else
                    {
                        file->flags |= FLAG_MISC;
                    }
                }
            }
            else
            {
                if (S_ISREG(file->st.st_mode))
                    s->num_reg++;
                else if (S_ISDIR(file->st.st_mode))
                    s->num_dir++;
                else if (S_ISCHR(file->st.st_mode))
                    s->num_chr++;
                else if (S_ISBLK(file->st.st_mode))
                    s->num_blk++;
                else if (S_ISFIFO(file->st.st_mode))
                    s->num_pipe++;
                else if (S_ISSOCK(file->st.st_mode))
                    s->num_sock++;
                else
                    file->flags |= FLAG_MISC;
            }

            if (file->flags & FLAG_MISC)
                s->num_misc++;

            if (!snapshot_buf_avail(s, s->paths_buf_used, readlink_len,
                                    &(s->paths_buf_size), INITIAL_NAMES,
                                    "file path", &(s->paths_buf),
                                    fix_file_paths))
            {
                return 0;
            }

            file->path = s->paths_buf + s->paths_buf_used;
            memcpy(file->path, s->readlink_buf, readlink_len);
            num_str = strrchr(fd_glob.gl_pathv[j], '/') + 1;
            file->fd = strtoul(num_str, &num_end, 0);
            if ((num_end == num_str) || (*num_end != '\0'))
            {
                snapshot_log(s, LOG_ERR,
                             "PID %u file %s is not a file descriptor",
                             proc->pid, fd_glob.gl_pathv[j]);

                continue;
            }

            snprintf(tmp_path, sizeof(tmp_path), "/proc/%u/fdinfo/%u",
                     proc->pid, file->fd);

            if (!snapshot_read_file(s, tmp_path, &read_bytes))
                continue;

            if (strncmp(s->read_buf, "pos:\t", 5) != 0)
            {
                snapshot_log(s, LOG_ERR, "'pos:' not where expected in %s",
                             tmp_path);

                continue;
            }

            num_str = s->read_buf + 5;
            file->pos = strtoull(num_str, &num_end, 0);
            if ((num_end == num_str) || (*num_end != '\n'))
            {
                snapshot_log(s, LOG_ERR, "position not numeric in %s",
                             tmp_path);

                continue;
            }

            if (strncmp(num_end, "\nflags:\t", 8) != 0)
            {
                snapshot_log(s, LOG_ERR, "'flags:' not where expected in %s",
                             tmp_path);

                continue;
            }

            num_str = num_end + 8;
            file->fdinfo_flags = strtoul(num_str, &num_end, 8);
            if ((num_end == num_str) || (*num_end != '\n'))
            {
                snapshot_log(s, LOG_ERR, "flags not octal in %s", tmp_path);
                continue;
            }

            s->paths_buf_used += readlink_len;

            if (proc->num_files++ == 0)
                proc->file = file;

            s->num_files++;
        }

        if (proc->num_files == 0)
        {
            snapshot_log(s, LOG_INFO,
                         "no interesting files in PID %u (%u fds)", proc->pid,
                         fd_glob.gl_pathc);

            continue;
        }

        snprintf(tmp_path, sizeof(tmp_path), "/proc/%u/", proc->pid);
        if (stat(tmp_path, &st) < 0)
        {
            snapshot_log(s, LOG_ERR,
                         "skipping PID %u (%s on %s)", proc->pid,
                         strerror(errno), tmp_path);

            continue;
        }

        proc->uid = st.st_uid;

        snprintf(tmp_path, sizeof(tmp_path), "/proc/%u/cmdline", proc->pid);
        if (stat(tmp_path, &st) < 0)
        {
            snapshot_log(s, LOG_ERR,
                         "skipping PID %u (%s on %s)", proc->pid,
                         strerror(errno), tmp_path);

            continue;
        }

        if (!snapshot_read_file(s, tmp_path, &read_bytes))
            continue;

        /* TO-DO: add '' around arguments with spaces */
        if (!snapshot_buf_avail(s, s->cmds_buf_used, read_bytes,
                                &(s->cmds_buf_size), INITIAL_CMDS,
                                "command line", &(s->cmds_buf), fix_proc_cmds))
        {
            return 0;
        }

        proc->cmdline = s->cmds_buf + s->cmds_buf_used;
        for (j = 0; j <= read_bytes; j++)
        {
            if ((s->read_buf[j] == '\0') && (j < (read_bytes - 1)))
                proc->cmdline[j] = ' ';
            else
                proc->cmdline[j] = s->read_buf[j];
        }

        s->cmds_buf_used += read_bytes;
        s->num_processes++;
        globfree(&fd_glob);
    }

    globfree(&proc_glob);
    return 1;
}

static int snapshot_read_file(snapshot *s, char *file, size_t *bytes)
{
    int fd;
    ssize_t cur_bytes;
    size_t total_bytes;
    char *old_read_buf;

    if (s->read_buf == NULL)
    {
        s->read_buf_size = INITIAL_READ;
        if ((s->read_buf = malloc(s->read_buf_size)) == NULL)
        {
            snapshot_log(s, LOG_ERR, "could not allocate file input buffer");
            return 0;
        }
    }

    if ((fd = open(file, O_RDONLY)) < 0)
    {
        snapshot_log(s, LOG_ERR, "could not open file %s (%s)", file,
                     strerror(errno));

        return 0;
    }

    total_bytes = 0;
    do
    {
        cur_bytes = read(fd, s->read_buf + total_bytes,
                         s->read_buf_size - total_bytes);

        if (cur_bytes < 0)
        {
            snapshot_log(s, LOG_ERR, "could not read file %s (%s)", file,
                         strerror(errno));

            close(fd);
            return 0;
        }

        total_bytes += cur_bytes;
        if (total_bytes == s->read_buf_size)
        {
            snapshot_log(s, LOG_DEBUG,
                         "growing file input buffer from %u to %u bytes",
                         s->read_buf_size, s->read_buf_size << 1);

            old_read_buf = s->read_buf;
            s->read_buf_size <<= 1;
            if ((s->read_buf = malloc(s->read_buf_size)) == NULL)
            {
                snapshot_log(s, LOG_ERR,
                             "could not allocate file input buffer");

                close(fd);
                return 0;
            }

            memcpy(s->read_buf, old_read_buf, total_bytes);
            free(old_read_buf);
        }
    } while (cur_bytes != 0);

    close(fd);
    s->read_buf[total_bytes] = '\0';
    *bytes = total_bytes;

    return 1;
}

static int snapshot_buf_avail(snapshot *s, size_t buf_used, size_t bytes,
                              size_t *buf_size, size_t initial, char *name,
                              void *vbuf, snapshot_fixup_function f)
{
    char *old_buf, *new_buf;
    char **buf;
    size_t new_size;

    if ((buf_used + bytes) < *buf_size)
        return 1;

    buf = vbuf;
    old_buf = *buf;
    if (*buf_size == 0)
        new_size = initial;
    else
        new_size = *buf_size;

    while ((buf_used + bytes) >= new_size)
        new_size <<= 1;

    snapshot_log(s, LOG_DEBUG, "growing %s buffer from %u to %u bytes", name,
                 *buf_size, new_size);

    if ((new_buf = malloc(new_size)) == NULL)
    {
        snapshot_log(s, LOG_ERR, "could not allocate %s buffer", name);
        return 0;
    }

    memset(new_buf, 0, new_size);
    if (old_buf != NULL)
    {
        memcpy(new_buf, old_buf, buf_used);
        f(s, new_buf);
        free(old_buf);
    }

    *buf_size = new_size;
    *buf = new_buf;
    return 1;
}

static int fix_file_paths(snapshot *s, char *new_buf)
{
    size_t i;

    for (i = 0; i < s->num_files; i++)
        s->file[i].path = new_buf + (s->file[i].path - s->paths_buf);

    return 1;
}

static int fix_proc_cmds(snapshot *s, char *new_buf)
{
    size_t i;

    for (i = 0; i < s->num_processes; i++)
        s->proc[i].cmdline = new_buf + (s->proc[i].cmdline - s->cmds_buf);

    return 1;
}

static int fix_proc_file(snapshot *s, char *new_buf)
{
    size_t i;
    file_info *new_file = (file_info *)new_buf;

    for (i = 0; i <= s->num_processes; i++)
        s->proc[i].file = new_file + (s->proc[i].file - s->file);

    return 1;
}
