/*
 * CDDL HEADER START
 *
 * The contents of this file are subject to the terms of the
 * Common Development and Distribution License (the "License").
 * You may not use this file except in compliance with the License.
 *
 * You can obtain a copy of the license at usr/src/OPENSOLARIS.LICENSE
 * or http://www.opensolaris.org/os/licensing.
 * See the License for the specific language governing permissions
 * and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL HEADER in each
 * file and include the License file at usr/src/OPENSOLARIS.LICENSE.
 * If applicable, add the following below this CDDL HEADER, with the
 * fields enclosed by brackets "[]" replaced with your own identifying
 * information: Portions Copyright [yyyy] [name of copyright owner]
 *
 * CDDL HEADER END
 */
/*
 * Copyright 2006 Sun Microsystems, Inc.  All rights reserved.
 * Use is subject to license terms.
 */
/* Copyright (c) 2012 by Delphix. All rights reserved */

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/param.h>
#include <sys/task.h>
#include <sys/contract.h>

#include <signal.h>
#include <unistd.h>
#include <dirent.h>
#include <stdlib.h>
#include <string.h>

#include <libintl.h>
#include <locale.h>
#include <stdio.h>
#include <fcntl.h>
#include <ctype.h>
#include <wchar.h>
#include <limits.h>
#include <libuutil.h>
#include <libcontract_priv.h>

#include <procfs.h>
#include <project.h>
#include <pwd.h>
#include <grp.h>
#include <zone.h>

#include "psexp.h"
#include "pgrep.h"

#ifndef TEXT_DOMAIN
#define TEXT_DOMAIN     "SYS_TEST"
#endif

#define OPT_SETB        0x0001  /* Set the bits specified by o_bits */
#define OPT_CLRB        0x0002  /* Clear the bits specified by o_bits */
#define OPT_FUNC        0x0004  /* Call the function specified by o_func */
#define OPT_STR         0x0008  /* Set the string specified by o_ptr */
#define OPT_CRIT        0x0010  /* Option is part of selection criteria */

#define F_LONG_FMT      0x0001  /* Match against long format cmd */
#define F_NEWEST        0x0002  /* Match only newest pid */
#define F_REVERSE       0x0004  /* Reverse matching criteria */
#define F_EXACT_MATCH   0x0008  /* Require exact match */
#define F_HAVE_CRIT     0x0010  /* Criteria specified */
#define F_OUTPUT        0x0020  /* Some output has been printed */
#define F_KILL          0x0040  /* Pkill semantics active (vs pgrep) */
#define F_LONG_OUT      0x0080  /* Long output format (pgrep -l) */
#define F_OLDEST        0x0100  /* Match only oldest pid */

static int opt_euid(char, char *);
static int opt_uid(char, char *);
static int opt_gid(char, char *);
static int opt_ppid(char, char *);
static int opt_pgrp(char, char *);
static int opt_sid(char, char *);
static int opt_term(char, char *);
static int opt_projid(char, char *);
static int opt_taskid(char, char *);
static int opt_zoneid(char, char *);
static int opt_ctid(char, char *);

static const char *g_procdir = "/proc"; /* Default procfs mount point */
static const char *g_delim = "\n";      /* Default output delimiter */
static const char *g_pname;             /* Program name for error messages */
static ushort_t g_flags;                /* Miscellaneous flags */

static optdesc_t g_optdtab[] = {
        { 0, 0, 0, 0 },                                 /* 'A' */
        { 0, 0, 0, 0 },                                 /* 'B' */
        { 0, 0, 0, 0 },                                 /* 'C' */
        { OPT_STR, 0, 0, &g_procdir },                  /* -D procfsdir */
        { 0, 0, 0, 0 },                                 /* 'E' */
        { 0, 0, 0, 0 },                                 /* 'F' */
        { OPT_FUNC | OPT_CRIT, 0, opt_gid, 0 },         /* -G gid */
        { 0, 0, 0, 0 },                                 /* 'H' */
        { 0, 0, 0, 0 },                                 /* 'I' */
        { OPT_FUNC | OPT_CRIT, 0, opt_projid, 0 },      /* -J projid */
        { 0, 0, 0, 0 },                                 /* 'K' */
        { 0, 0, 0, 0 },                                 /* 'L' */
        { 0, 0, 0, 0 },                                 /* 'M' */
        { 0, 0, 0, 0 },                                 /* 'N' */
        { 0, 0, 0, 0 },                                 /* 'O' */
        { OPT_FUNC | OPT_CRIT, 0, opt_ppid, 0 },        /* -P ppid */
        { 0, 0, 0, 0 },                                 /* 'Q' */
        { 0, 0, 0, 0 },                                 /* 'R' */
        { 0, 0, 0, 0 },                                 /* 'S' */
        { OPT_FUNC | OPT_CRIT, 0, opt_taskid, 0 },      /* -T taskid */
        { OPT_FUNC | OPT_CRIT, 0, opt_uid, 0 },         /* -U uid */
        { 0, 0, 0, 0 },                                 /* 'V' */
        { 0, 0, 0, 0 },                                 /* 'W' */
        { 0, 0, 0, 0 },                                 /* 'X' */
        { 0, 0, 0, 0 },                                 /* 'Y' */
        { 0, 0, 0, 0 },                                 /* 'Z' */
        { 0, 0, 0, 0 },                                 /* '[' */
        { 0, 0, 0, 0 },                                 /* '\\' */
        { 0, 0, 0, 0 },                                 /* ']' */
        { 0, 0, 0, 0 },                                 /* '^' */
        { 0, 0, 0, 0 },                                 /* '_' */
        { 0, 0, 0, 0 },                                 /* '`' */
        { 0, 0, 0, 0 },                                 /* 'a' */
        { 0, 0, 0, 0 },                                 /* 'b' */
        { OPT_FUNC | OPT_CRIT, 0, opt_ctid, 0 },        /* -c ctid */
        { OPT_STR, 0, 0, &g_delim },                    /* -d delim */
        { 0, 0, 0, 0 },                                 /* 'e' */
        { OPT_SETB, F_LONG_FMT, 0, &g_flags },          /* -f */
        { OPT_FUNC | OPT_CRIT, 0, opt_pgrp, 0 },        /* -g pgrp */
        { 0, 0, 0, 0 },                                 /* 'h' */
        { 0, 0, 0, 0 },                                 /* 'i' */
        { 0, 0, 0, 0 },                                 /* 'j' */
        { 0, 0, 0, 0 },                                 /* 'k' */
        { OPT_SETB, F_LONG_OUT, 0, &g_flags },          /* 'l' */
        { 0, 0, 0, 0 },                                 /* 'm' */
        { OPT_SETB, F_NEWEST, 0, &g_flags },            /* -n */
        { OPT_SETB, F_OLDEST, 0, &g_flags },            /* -o */
        { 0, 0, 0, 0 },                                 /* 'p' */
        { 0, 0, 0, 0 },                                 /* 'q' */
        { 0, 0, 0, 0 },                                 /* 'r' */
        { OPT_FUNC | OPT_CRIT, 0, opt_sid, 0 },         /* -s sid */
        { OPT_FUNC | OPT_CRIT, 0, opt_term, 0 },        /* -t term */
        { OPT_FUNC | OPT_CRIT, 0, opt_euid, 0 },        /* -u euid */
        { OPT_SETB, F_REVERSE, 0, &g_flags },           /* -v */
        { 0, 0, 0, 0 },                                 /* 'w' */
        { OPT_SETB, F_EXACT_MATCH, 0, &g_flags },       /* -x */
        { 0, 0, 0, 0 },                                 /* 'y' */
        { OPT_FUNC | OPT_CRIT, 0, opt_zoneid, 0 }       /* -z zoneid */
};

static const char PGREP_USAGE[] = "\
Usage: %s [-flnovx] [-d delim] [-P ppidlist] [-g pgrplist] [-s sidlist]\n\
        [-u euidlist] [-U uidlist] [-G gidlist] [-J projidlist]\n\
        [-T taskidlist] [-t termlist] [-z zonelist] [-c ctidlist] [pattern]\n";

static const char PKILL_USAGE[] = "\
Usage: %s [-signal] [-fnovx] [-P ppidlist] [-g pgrplist] [-s sidlist]\n\
        [-u euidlist] [-U uidlist] [-G gidlist] [-J projidlist]\n\
        [-T taskidlist] [-t termlist] [-z zonelist] [-c ctidlist] [pattern]\n";

static const char PGREP_OPTS[] = ":flnovxc:d:D:u:U:G:P:g:s:t:z:J:T:";
static const char PKILL_OPTS[] = ":fnovxc:D:u:U:G:P:g:s:t:z:J:T:";

static const char LSEP[] = ",\t ";      /* Argument list delimiter chars */

static psexp_t g_psexp;                 /* Process matching expression */
static pid_t g_pid;                     /* Current pid */
static int g_signal = SIGTERM;          /* Signal to send */

static void
print_proc(psinfo_t *psinfo)
{
        if (g_flags & F_OUTPUT)
                (void) printf("%s%d", g_delim, (int)psinfo->pr_pid);
        else {
                (void) printf("%d", (int)psinfo->pr_pid);
                g_flags |= F_OUTPUT;
        }
}

static char *
mbstrip(char *buf, size_t nbytes)
{
        wchar_t wc;
        char *p;
        int n;

        buf[nbytes - 1] = '\0';
        p = buf;

        while (*p != '\0') {
                n = mbtowc(&wc, p, MB_LEN_MAX);

                if (n < 0 || !iswprint(wc)) {
                        if (n < 0)
                                n = sizeof (char);

                        if (nbytes <= n) {
                                *p = '\0';
                                break;
                        }

                        (void) memmove(p, p + n, nbytes - n);

                } else {
                        nbytes -= n;
                        p += n;
                }
        }

        return (buf);
}

static void
print_proc_long(psinfo_t *psinfo)
{
        char *name;

        if (g_flags & F_LONG_FMT)
                name = mbstrip(psinfo->pr_psargs, PRARGSZ);
        else
                name = psinfo->pr_fname;

        if (g_flags & F_OUTPUT)
                (void) printf("%s%5d %s", g_delim, (int)psinfo->pr_pid, name);
        else {
                (void) printf("%5d %s", (int)psinfo->pr_pid, name);
                g_flags |= F_OUTPUT;
        }
}

static void
kill_proc(psinfo_t *psinfo)
{
        if (psinfo->pr_pid > 0 && kill(psinfo->pr_pid, g_signal) == -1)
                uu_warn(gettext("Failed to signal pid %d"),
                    (int)psinfo->pr_pid);
}

static DIR *
open_proc_dir(const char *dirpath)
{
        struct stat buf;
        DIR *dirp;

        if ((dirp = opendir(dirpath)) == NULL) {
                uu_warn(gettext("Failed to open %s"), dirpath);
                return (NULL);
        }

        if (fstat(dirp->dd_fd, &buf) == -1) {
                uu_warn(gettext("Failed to stat %s"), dirpath);
                (void) closedir(dirp);
                return (NULL);
        }

        if (strcmp(buf.st_fstype, "proc") != 0) {
                uu_warn(gettext("%s is not a procfs mount point\n"), dirpath);
                (void) closedir(dirp);
                return (NULL);
        }

        return (dirp);
}

#define NEWER(ps1, ps2) \
        ((ps1.pr_start.tv_sec > ps2.pr_start.tv_sec) || \
            (ps1.pr_start.tv_sec == ps2.pr_start.tv_sec && \
            ps1.pr_start.tv_nsec > ps2.pr_start.tv_nsec))

static int
scan_proc_dir(const char *dirpath, DIR *dirp, psexp_t *psexp,
        void (*funcp)(psinfo_t *))
{
        char procpath[MAXPATHLEN];
        psinfo_t ps, ops;
        dirent_t *dent;
        int procfd;

        int reverse = (g_flags & F_REVERSE) ? 1 : 0;
        int ovalid = 0, nmatches = 0, flags = 0;

        if (g_flags & F_LONG_FMT)
                flags |= PSEXP_PSARGS;

        if (g_flags & F_EXACT_MATCH)
                flags |= PSEXP_EXACT;

        while ((dent = readdir(dirp)) != NULL) {

                if (dent->d_name[0] == '.')
                        continue;

                (void) snprintf(procpath, sizeof (procpath), "%s/%s/psinfo",
                    dirpath, dent->d_name);

                if ((procfd = open(procpath, O_RDONLY)) == -1)
                        continue;

                if ((read(procfd, &ps, sizeof (ps)) == sizeof (psinfo_t)) &&
                    (ps.pr_nlwp != 0) && (ps.pr_pid != g_pid) &&
                    (psexp_match(psexp, &ps, flags) ^ reverse)) {

                        if (g_flags & F_NEWEST) {
                                /* LINTED - opsinfo use ok */
                                if (!ovalid || NEWER(ps, ops)) {
                                        (void) memcpy(&ops, &ps,
                                            sizeof (psinfo_t));
                                        ovalid = 1;
                                }
                        } else if (g_flags & F_OLDEST) {
                                if (!ovalid || NEWER(ops, ps)) {
                                        (void) memcpy(&ops, &ps,
                                            sizeof (psinfo_t));
                                        ovalid = 1;
                                }
                        } else {
                                (*funcp)(&ps);
                                nmatches++;
                        }
                }

                (void) close(procfd);
        }

        if ((g_flags & (F_NEWEST | F_OLDEST)) && ovalid) {
                (*funcp)(&ops);
                nmatches++;
        }

        return (nmatches);
}

static int
parse_ids(idtab_t *idt, char *arg, int base, int opt, idkey_t zero)
{
        char *ptr, *next;
        idkey_t id;

        for (ptr = strtok(arg, LSEP); ptr != NULL; ptr = strtok(NULL, LSEP)) {
                if ((id = (idkey_t)strtoul(ptr, &next, base)) != 0)
                        idtab_append(idt, id);
                else
                        idtab_append(idt, zero);

                if (next == ptr || *next != 0) {
                        uu_warn("invalid argument for option '%c' -- %s\n",
                            opt, ptr);
                        return (-1);
                }
        }

        return (0);
}

static int
parse_uids(idtab_t *idt, char *arg)
{
        char *ptr, *next;
        struct passwd *pwent;
        idkey_t id;

        for (ptr = strtok(arg, LSEP); ptr != NULL; ptr = strtok(NULL, LSEP)) {
                if (isdigit(ptr[0])) {
                        id = strtol(ptr, &next, 10);

                        if (next != ptr && *next == '\0') {
                                idtab_append(idt, id);
                                continue;
                        }
                }

                if ((pwent = getpwnam(ptr)) != NULL)
                        idtab_append(idt, pwent->pw_uid);
                else
                        goto err;
        }

        return (0);

err:
        uu_warn(gettext("invalid user name -- %s\n"), ptr);
        return (-1);
}

static int
parse_gids(idtab_t *idt, char *arg)
{
        char *ptr, *next;
        struct group *grent;
        idkey_t id;

        for (ptr = strtok(arg, LSEP); ptr != NULL; ptr = strtok(NULL, LSEP)) {
                if (isdigit(ptr[0])) {
                        id = strtol(ptr, &next, 10);

                        if (next != ptr && *next == '\0') {
                                idtab_append(idt, id);
                                continue;
                        }
                }

                if ((grent = getgrnam(ptr)) != NULL)
                        idtab_append(idt, grent->gr_gid);
                else
                        goto err;
        }

        return (0);

err:
        uu_warn(gettext("invalid group name -- %s\n"), ptr);
        return (-1);
}

static int
parse_ttys(idtab_t *idt, char *arg)
{
        char devpath[MAXPATHLEN];
        struct stat buf;
        char *ptr;

        int seen_console = 0; /* Flag so we only stat syscon and systty once */

        for (ptr = strtok(arg, LSEP); ptr != NULL; ptr = strtok(NULL, LSEP)) {
                if (strcmp(ptr, "none") == 0) {
                        idtab_append(idt, (idkey_t)PRNODEV);
                        continue;
                }

                if (strcmp(ptr, "console") == 0) {
                        if (seen_console)
                                continue;

                        if (stat("/dev/syscon", &buf) == 0)
                                idtab_append(idt, (idkey_t)buf.st_rdev);

                        if (stat("/dev/systty", &buf) == 0)
                                idtab_append(idt, (idkey_t)buf.st_rdev);

                        seen_console++;
                }

                (void) snprintf(devpath, MAXPATHLEN - 1, "/dev/%s", ptr);

                if (stat(devpath, &buf) == -1)
                        goto err;

                idtab_append(idt, (idkey_t)buf.st_rdev);
        }

        return (0);

err:
        uu_warn(gettext("unknown terminal name -- %s\n"), ptr);
        return (-1);
}

static int
parse_projects(idtab_t *idt, char *arg)
{
        char *ptr, *next;
        projid_t projid;
        idkey_t id;

        for (ptr = strtok(arg, LSEP); ptr != NULL; ptr = strtok(NULL, LSEP)) {
                if (isdigit(ptr[0])) {
                        id = strtol(ptr, &next, 10);

                        if (next != ptr && *next == '\0') {
                                idtab_append(idt, id);
                                continue;
                        }
                }

                if ((projid = getprojidbyname(ptr)) != -1)
                        idtab_append(idt, projid);
                else
                        goto err;
        }

        return (0);

err:
        uu_warn(gettext("invalid project name -- %s\n"), ptr);
        return (-1);
}

static int
parse_zones(idtab_t *idt, char *arg)
{
        char *ptr;
        zoneid_t id;

        for (ptr = strtok(arg, LSEP); ptr != NULL; ptr = strtok(NULL, LSEP)) {
                if (zone_get_id(ptr, &id) != 0) {
                        uu_warn(gettext("invalid zone name -- %s\n"), ptr);
                        return (-1);
                }
                idtab_append(idt, id);
        }

        return (0);
}

/*ARGSUSED*/
static int
opt_euid(char c, char *arg)
{
        return (parse_uids(&g_psexp.ps_euids, arg));
}

/*ARGSUSED*/
static int
opt_uid(char c, char *arg)
{
        return (parse_uids(&g_psexp.ps_ruids, arg));
}

/*ARGSUSED*/
static int
opt_gid(char c, char *arg)
{
        return (parse_gids(&g_psexp.ps_rgids, arg));
}

static int
opt_ppid(char c, char *arg)
{
        return (parse_ids(&g_psexp.ps_ppids, arg, 10, c, 0));
}

static int
opt_pgrp(char c, char *arg)
{
        return (parse_ids(&g_psexp.ps_pgids, arg, 10, c, getpgrp()));
}

static int
opt_sid(char c, char *arg)
{
        return (parse_ids(&g_psexp.ps_sids, arg, 10, c, getsid(0)));
}

/*ARGSUSED*/
static int
opt_term(char c, char *arg)
{
        return (parse_ttys(&g_psexp.ps_ttys, arg));
}

/*ARGSUSED*/
static int
opt_projid(char c, char *arg)
{
        return (parse_projects(&g_psexp.ps_projids, arg));
}

static int
opt_taskid(char c, char *arg)
{
        return (parse_ids(&g_psexp.ps_taskids, arg, 10, c, gettaskid()));
}

/*ARGSUSED*/
static int
opt_zoneid(char c, char *arg)
{
        return (parse_zones(&g_psexp.ps_zoneids, arg));
}

static int
opt_ctid(char c, char *arg)
{
        return (parse_ids(&g_psexp.ps_ctids, arg, 10, c, getctid()));
}

static void
print_usage(FILE *stream)
{
        if (g_flags & F_KILL)
                (void) fprintf(stream, gettext(PKILL_USAGE), g_pname);
        else
                (void) fprintf(stream, gettext(PGREP_USAGE), g_pname);
}

int
main(int argc, char *argv[])
{
        void (*funcp)(psinfo_t *);

        const char *optstr;
        optdesc_t *optd;
        int nmatches, c;

        DIR *dirp;

        (void) setlocale(LC_ALL, "");
        (void) textdomain(TEXT_DOMAIN);

        UU_EXIT_FATAL = E_ERROR;

        g_pname = uu_setpname(argv[0]);
        g_pid = getpid();

        psexp_create(&g_psexp);

        if (strcmp(g_pname, "pkill") == 0) {

                if (argc > 1 && argv[1][0] == '-' &&
                    str2sig(&argv[1][1], &g_signal) == 0) {
                        argv[1] = argv[0];
                        argv++;
                        argc--;
                }

                optstr = PKILL_OPTS;
                g_flags |= F_KILL;
        } else
                optstr = PGREP_OPTS;

        opterr = 0;

        while (optind < argc) {
                while ((c = getopt(argc, argv, optstr)) != (int)EOF) {

                        if (c == ':' || c == '?' ||
                            g_optdtab[c - 'A'].o_opts == 0) {
                                if (c == ':') {
                                        uu_warn(
                                            gettext("missing argument -- %c\n"),
                                            optopt);
                                } else if (optopt != '?') {
                                        uu_warn(
                                            gettext("illegal option -- %c\n"),
                                            optopt);
                                }

                                print_usage(stderr);
                                return (E_USAGE);
                        }

                        optd = &g_optdtab[c - 'A'];

                        if (optd->o_opts & OPT_SETB)
                                *((ushort_t *)optd->o_ptr) |= optd->o_bits;

                        if (optd->o_opts & OPT_CLRB)
                                *((ushort_t *)optd->o_ptr) &= ~optd->o_bits;

                        if (optd->o_opts & OPT_STR)
                                *((char **)optd->o_ptr) = optarg;

                        if (optd->o_opts & OPT_CRIT)
                                g_flags |= F_HAVE_CRIT;

                        if (optd->o_opts & OPT_FUNC) {
                                if (optd->o_func(c, optarg) == -1)
                                        return (E_USAGE);
                        }
                }

                if (optind < argc) {
                        if (g_psexp.ps_pat != NULL) {
                                uu_warn(gettext("illegal argument -- %s\n"),
                                    argv[optind]);
                                print_usage(stderr);
                                return (E_USAGE);
                        }

                        g_psexp.ps_pat = argv[optind++];
                        g_flags |= F_HAVE_CRIT;
                }
        }

        if ((g_flags & F_NEWEST) && (g_flags & F_OLDEST)) {
                uu_warn(gettext("-n and -o are mutually exclusive\n"));
                print_usage(stderr);
                return (E_USAGE);
        }

        if ((g_flags & F_HAVE_CRIT) == 0) {
                uu_warn(gettext("No matching criteria specified\n"));
                print_usage(stderr);
                return (E_USAGE);
        }

        if (psexp_compile(&g_psexp) == -1) {
                psexp_destroy(&g_psexp);
                return (E_USAGE);
        }

        if ((dirp = open_proc_dir(g_procdir)) == NULL)
                return (E_ERROR);

        if (g_flags & F_KILL)
                funcp = kill_proc;
        else if (g_flags & F_LONG_OUT)
                funcp = print_proc_long;
        else
                funcp = print_proc;

        nmatches = scan_proc_dir(g_procdir, dirp, &g_psexp, funcp);

        if (g_flags & F_OUTPUT)
                (void) fputc('\n', stdout);

        psexp_destroy(&g_psexp);
        return (nmatches ? E_MATCH : E_NOMATCH);
}