#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <stdarg.h>
#include <fcntl.h>
#include <signal.h>
#include <pwd.h>
#include <grp.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <sys/prctl.h>
#include <assert.h>

#include "utils.h"
#include "log.h"

#define _POSIX_PATH_MAX 256

char *arg0 = NULL;
static int null_fd = -1;


static void sighandler_child(int signo)
{
    switch (signo) {
        case SIGHUP:
            if (getppid() == 1) {
                N("Master process %d died, exiting", getpgrp());
                exit(EXIT_SUCCESS);
            }
            break;
    }
}

int set_child_sighandler(void)
{
    if (prctl(PR_SET_PDEATHSIG, SIGHUP) != 0)
        return 1;
    assert( signal(SIGCHLD, SIG_IGN) != SIG_ERR );
    assert( signal(SIGPIPE, SIG_IGN) != SIG_ERR );

    return signal(SIGHUP, sighandler_child) == SIG_ERR;
}

void set_procname(const char *new_arg0)
{
    assert(arg0);
	memset(arg0, 0, _POSIX_PATH_MAX);
	strncpy(arg0, new_arg0, _POSIX_PATH_MAX);
}

pid_t daemonize(int stay_foreground)
{
    int status = -1;
    pid_t pid;

    /* Fork off the parent process */
    pid = fork();

    /* An error occurred */
    if (pid < 0)
        return pid;

    /* Success: Let the parent terminate */
    if (pid > 0) {
        if (!stay_foreground)
            exit(EXIT_SUCCESS);
        waitpid(-1, &status, 0);
        exit(EXIT_SUCCESS);
    }

    /* On success: The child process becomes session leader */
    if (!stay_foreground && setsid() < 0) {
        E_STRERR("setsid");
        exit(EXIT_FAILURE);
    }

    /* Fork off for the second time*/
    if (!stay_foreground) {
        pid = fork();

        /* An error occurred */
        if (pid < 0)
            exit(EXIT_FAILURE);

        /* Success: Let the parent terminate */
        if (pid > 0)
            exit(EXIT_SUCCESS);
    }

    if (!stay_foreground && setpgrp()) {
        E_STRERR("setpgrp");
        exit(EXIT_FAILURE);
    }

    /* Set new file permissions */
    umask(0);

    if (!stay_foreground) {
        /* Change the working directory to the root directory */
        /* or another appropriated directory */
        chdir("/");
        /* Close all open file descriptors */
        assert( close_fds_except(-1) == 0 );
        assert( redirect_devnull_to(0, 1, 2, -1) == 0 );
    } else {
        assert( close_fds_except(0, 1, 2, -1) == 0 );
    }

    return pid;
}

int close_fds_except(int fds, ...)
{
    int fd;
    long max_fd;
    size_t i, except_count, found;
    va_list ap;

    max_fd = sysconf(_SC_OPEN_MAX);
    if (max_fd <= 0)
        return 1;

    va_start(ap, fds);
    {
        int *all_fds = malloc(max_fd * sizeof(*all_fds));
        assert(all_fds);
        memset(all_fds, -1, max_fd * sizeof(*all_fds));

        except_count = 0;
        while ( (fd = va_arg(ap, int)) >= 0 ) {
            all_fds[except_count++] = fd;
        }

        for (fd = max_fd; fd >= 0; --fd) {
            found = 0;
            for (i = 0; i < except_count; ++i) {
                if (fd == all_fds[i])
                    found++;
            }
            if (!found) {
                close(fd);
            }
        }

        free(all_fds);
    }
    va_end(ap);

    return 0;
}

int redirect_devnull_to(int fds, ...)
{
    int fd, rc = 0;
    va_list ap;

    if (null_fd < 0)
        null_fd = open("/dev/null", O_RDWR);
    if (null_fd < 0)
        return -1;
    assert(null_fd >= 0);

    va_start(ap, fds);
    {
        while ( (fd = va_arg(ap, int)) >= 0 ) {
            if ( dup2(null_fd, fd) < 0 )
                rc++;
        }
    }
    va_end(ap);

    return rc;
}

int change_user_group(const char *user, const char *group)
{
    struct passwd *pwd = NULL;
    struct group *grp = NULL;
    gid_t gid;

    pwd = getpwnam(user);
    if (!pwd)
        return 1;

    if (!group) {
        gid = pwd->pw_gid;
    } else {
        grp = getgrnam(group);
        if (!grp)
            return 1;
        gid = grp->gr_gid;
    }

    if (setregid(gid, gid))
        return 1;
    if (setreuid(pwd->pw_uid, pwd->pw_uid))
        return 1;

    return 0;
}