/*
 * QUOTA    An implementation of the diskquota system for the LINUX operating
 *          system. QUOTA is implemented using the BSD systemcall interface
 *          as the means of communication with the user level. Should work for
 *          all filesystems because of integration into the VFS layer of the
 *          operating system. This is based on the Melbourne quota system wich
 *          uses both user and group quota files.
 * 
 *          Program to check disk quotas.
 * 
 * Authors:
 *          Directory reading routines: Edvard Tuinder <ed@delirium.nl.mugnet.org>
 *          Quota storing routines: Marco van Wieringen <mvw@mcs.nl.mugnet.org>
 *
 *          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 2 of
 *          the License, or (at your option) any later version.
 */

#include <sys/types.h>
#include <sys/param.h>
#include <dirent.h>
#include <sys/stat.h>
#include <stdio.h>
#include <linux/quota.h>
#include <stdarg.h>
#include <sys/file.h>
#include <mntent.h>
#include <getopt.h>
#include <limits.h>
#include <unistd.h>
#include <stdlib.h>

#define DEF_BLOCKSIZE 1024
#define NODQUOT (struct dquot *)NULL

#ifndef lint
static char RCS_checkquota[] = "$Id: quotacheck.c,v 2.5 1994/02/09 19:00:50 mvw Exp mvw $";
#endif

struct dquot {
   int          dq_id;  /* id this applies to (uid, gid) */
   struct dqblk dq_dqb; /* diskquota for id */
   struct dquot *next;  /* pointer to next id */
};

struct dlinks {
   ino_t ino;
   size_t size;
   size_t blksize;
   u_long id;
   struct dlinks *next;
};

struct dirs {
   char *dir_name;
   struct dirs *next;
};

void add_to_quota(struct stat *, int);
void dump_to_file(char *, char *, int);
void store_dlinks(struct stat *, int);
void remove_list(void);
void scan_dir(char *);
void abort_now(const char *, const char *, ...);
void add_dlinks(u_long *, u_long *, int, int);

static char bits[] = "|/-\\";
#define BITS_SIZE 4 /* sizeof(bits) == 5 */

dev_t cur_dev;
char dflag = 0, vflag = 0;
char check_usr, check_grp;
long files_done, dirs_done;
u_long highestid[MAXQUOTAS];
char *quotatypes[] = INITQFNAMES;
char log_buf[16384]; /* for dflag stderr buffering */
#ifdef DEBUG_MALLOC
static size_t malloc_mem = 0;
static size_t free_mem = 0;
#endif

struct dquot *dquot_list[MAXQUOTAS] = {NODQUOT, NODQUOT};
struct dquot *mru_dquot[MAXQUOTAS] = {NODQUOT, NODQUOT};
struct dlinks *stored_links[MAXQUOTAS] = {(struct dlinks *)NULL, (struct dlinks *)NULL};

/*
 * This simple algorithm calculates the size of a file in blocks.
 * It is not perfect but works most of the time.
 */
static inline size_t isize_to_blocks(size_t isize, size_t blksize)
{
   u_long blocks;
   u_long indirect;

   if (!blksize)
      blksize = DEF_BLOCKSIZE;

   blocks = (isize / blksize) + ((isize % blksize) ? 1 : 0);
   if (blocks > 10) {
      indirect = ((blocks - 11) >> 8) + 1; /* single indirect blocks */
      if (blocks > (10 + 256)) {
         indirect += ((blocks - 267) >> 16) + 1; /* double indirect blocks */
         if (blocks > (10 + 256 + (256 << 8)))
            indirect++; /* triple indirect blocks */
      }
      blocks += indirect;
   }
   return blocks;
}

/*
 * Ok check each memory allocation.
 */
void *xmalloc(size_t size)
{
   void *ptr;

#ifdef DEBUG_MALLOC
   malloc_mem += size;
#endif
   ptr = malloc(size);
   if (ptr == (void *)NULL)
      abort_now("xmalloc", "Virtual memory exhausted\n");
   memset(ptr, 0, size);
   return(ptr);
}

/*
 * Do a lookup of a type of quota for a specific id. Use short cut with
 * most recently used dquot struct pointer.
 */
static inline struct dquot *lookup_dquot(int id, int type)
{
   register struct dquot *lptr = NODQUOT;

   /*
    * First fast lookup when same as used before.
    */
   if (mru_dquot[type] != NODQUOT && mru_dquot[type]->dq_id == id)
      return (mru_dquot[type]);

   for (lptr = dquot_list[type]; lptr != NODQUOT; lptr = lptr->next)
      if (lptr->dq_id == id) {
         mru_dquot[type] = lptr;
         return (lptr);
      }
   return(NODQUOT);
}

/*
 * Add a new dquot for a new id to the list.
 */
static inline struct dquot *add_dquot(int id, int type)
{
   register struct dquot *lptr;

   if (dflag)
      fprintf(stderr, "Adding dquot structure type %s for %d\n",
	      quotatypes[type], id);

   lptr = (struct dquot *)xmalloc(sizeof(struct dquot));

   lptr->dq_id = id;
   lptr->next = dquot_list[type];
   dquot_list[type] = lptr;
   lptr->dq_btime = lptr->dq_itime = (time_t) 0;

   if (id > highestid[type])
      highestid[type] = id;

   return(lptr);
}

/*
 * Check the programs arguments for a specific target.
 */
int oneof(char *target, char *list[], int cnt)
{
   register int i;

   for (i = 0; i < cnt; i++)
      if (strcmp(target, list[i]) == 0)
         return (i);
   return(-1);
}

/*
 * Show a blitting cursor as means of visual progress indicator.
 */
static void blit()
{
   static short bitc = 0;

   putc(bits[bitc], stdout);
   putc('\b', stdout);
   bitc++;
   bitc %= BITS_SIZE;
}

void usage()
{
   fputs("Usage:\n\tquotacheck [-g] [-u] [-vd] -a\n", stderr);
   fputs("\tquotacheck [-g] [-u] [-vd] filesys ...\n", stderr);
   exit(1);
}

int main(int argc, char **argv)
{
   FILE *fp;
   int cnt, ch;
   struct stat st;
   char aflag = 0, gflag = 0, uflag = 0;
   long argnum, done;
   char *usr_qfnp, *grp_qfnp;
   register struct mntent *mnt;

   while ((ch = getopt(argc, argv, "avugd")) != EOF) {
      switch (ch) {
         case 'a':
            aflag++;
            break;
         case 'g':
            gflag++;
            break;
         case 'u':
            uflag++;
            break;
         case 'd':
            dflag++;
	    setbuf(stderr, log_buf);
	    break;
         case 'v':
            vflag++;
	    setbuf(stdout, (char *)NULL);
            break;
         default:
            usage();
      }
   }
   argc -= optind;
   argv += optind;

   if (vflag && dflag)
     vflag = 0;

   if (!uflag && !gflag)
      uflag++;

   if (!aflag && argc == 0)
      usage();

   fp = setmntent(FSTAB, "r");
   while ((mnt = getmntent(fp)) != (struct mntent *) NULL) {
      check_usr = check_grp = 0;
      if (argc && ((argnum = oneof(mnt->mnt_dir, argv, argc)) >= 0) ||
                  ((argnum = oneof(mnt->mnt_fsname, argv, argc)) >= 0)) {
          done |= 1 << argnum;
      } else
         if (!aflag || hasmntopt(mnt, MNTOPT_NOAUTO) ||
                       hasmntopt(mnt, MNTOPT_NOQUOTA))
            continue;

      if (gflag && hasquota(mnt, GRPQUOTA, &grp_qfnp))
         check_grp++;
      if (uflag && hasquota(mnt, USRQUOTA, &usr_qfnp))
         check_usr++;
      if (check_usr || check_grp) {
	 if ((lstat(mnt->mnt_dir, &st)) == -1) {
	   fprintf(stderr, "%s: not found\n", mnt->mnt_dir);
	   perror("lstat");
           continue;
	 }

	 if (vflag)
	   fprintf(stdout,"Scanning %s [%s] ", mnt->mnt_fsname, mnt->mnt_dir);

	 if (S_ISDIR(st.st_mode)) {
           cur_dev = st.st_dev;
           files_done = dirs_done = 0;
           if (check_usr)
	      add_to_quota(&st, USRQUOTA);
           if (check_grp)
	      add_to_quota(&st, GRPQUOTA);
	   scan_dir(mnt->mnt_dir);
	   dirs_done++;
	   if (vflag)
	      fputs("done\n", stderr);
           if (vflag || dflag)
              fprintf(stderr, "Checked %d directories and %d files\n",
                      dirs_done, files_done);
	 } else {
	   fprintf(stderr, "%s: not a directory\n", mnt->mnt_dir);
	   exit(0);
	 }

         if (check_usr)
            dump_to_file(mnt->mnt_fsname, usr_qfnp, USRQUOTA);
         if (check_grp)
            dump_to_file(mnt->mnt_fsname, grp_qfnp, GRPQUOTA);
         remove_list();
      }
   }
   endmntent(fp);

   for (cnt = 0; cnt < argc; cnt++)
      if ((done & (1 << cnt)) == 0)
         fprintf(stderr, "%s not found in fstab\n", argv[cnt]);

#ifdef DEBUG_MALLOC
   fprintf (stderr, "Allocated %d bytes memory\nFree'd %d bytes\nLost %d bytes\n", malloc_mem,
	free_mem, malloc_mem - free_mem);
#endif
   exit(0);
}

/*
 * Scan a directory with the readdir systemcall. Stat the files and add the sizes
 * of the files to the appropriate quotas. When we find a dir we recursivly call
 * ourself to scan that dir.
 */
void scan_dir(char *pathname)
{
   struct dirs *dir_stack = {(struct dirs *)NULL};
   struct dirs *new_dir;
   struct dirent *de;
   struct stat st;
   DIR *dp;

   if ((dp = opendir(pathname)) == (DIR *)NULL)
      abort_now("opendir", "\n%s\n", pathname);

   chdir(pathname);
   while ((de = readdir(dp)) != (struct dirent *)NULL) {
      if (!strcmp(de->d_name, ".") || !strcmp(de->d_name, ".."))
         continue;
      if (vflag)
         blit();

      if ((lstat(de->d_name, &st)) == -1) {
	 fprintf(stderr, "Hmm, file `%s/%s' not found\n", pathname, de->d_name);
         fputs("Guess you'd better run fsck first !\nexiting...\n", stderr);
         perror("lstat");
         exit(1);
      }

      if (check_usr)
	 add_to_quota(&st, USRQUOTA);
      if (check_grp)
	 add_to_quota(&st, GRPQUOTA);

      if (S_ISDIR(st.st_mode)) {
         if (st.st_dev != cur_dev)
            continue;
         /*
          * Add this to the directory stack and check this later on.
          */
	 if (dflag)
	    fprintf(stderr, "pushd %s/%s\n", pathname, de->d_name);
         new_dir = xmalloc(sizeof(struct dirs));
         new_dir->dir_name = xmalloc(strlen(pathname) + strlen(de->d_name) + 2);
         sprintf(new_dir->dir_name, "%s/%s", pathname, de->d_name);
         new_dir->next = dir_stack;
         dir_stack = new_dir;
      } else {
	 if (dflag)
	    fprintf(stderr, "\tAdding %s size %d ino %d links %d\n", de->d_name,
		    st.st_size, st.st_ino, st.st_nlink);
         files_done++;
      }
   }
   closedir(dp);

   /*
    * Traverse the directory stack, and check it.
    */
   if (dflag)
      fputs("Scanning stored directories from directory stack\n", stderr);
   while (dir_stack != (struct dirs *)NULL) {
      new_dir = dir_stack;
      dir_stack = dir_stack->next;
      if (dflag)
         fprintf(stderr, "popd %s\nEntering directory %s\n",
                 new_dir->dir_name, new_dir->dir_name);
      scan_dir(new_dir->dir_name);
      dirs_done++;
#ifdef DEBUG_MALLOC
      free_mem += sizeof(struct dirs) + strlen(new_dir->dir_name) + 1;
#endif
      free(new_dir->dir_name);
      free(new_dir);
   }
   if (dflag)
      fprintf(stderr, "Leaving %s\n", pathname);
}
    
/*
 * Store a hardlinked file for later. Add the end we add this to a users
 * quota because we don't wanna count it more then ones.
 */
void store_dlinks(struct stat *st, int type)
{
   struct dlinks *lptr;

   if (dflag)
      fprintf(stderr, "Adding hardlink for ino %d\n", st->st_ino);

   for (lptr = stored_links[type]; lptr != (struct dlinks *)NULL; lptr = lptr->next)
      if (lptr->ino == st->st_ino)
         return;

   lptr = (struct dlinks *)xmalloc(sizeof(struct dlinks));

   if (type == USRQUOTA)
     lptr->id = st->st_uid;
   else 
     lptr->id = st->st_gid;

   lptr->ino = st->st_ino;
   lptr->size = st->st_size;
   lptr->blksize = st->st_blksize;

   lptr->next = stored_links[type];
   stored_links[type] = lptr;
}

/*
 * Add a number of blocks and inodes to a quota.
 */
void add_to_quota(struct stat *st, int type)
{
   int wanted;
   struct dquot *lptr;

   switch(type)
   {
      case USRQUOTA:
         wanted = st->st_uid;
         break;
      case GRPQUOTA:
         wanted = st->st_gid;
         break;
      default:
         return;
   }

   if ((lptr = lookup_dquot(wanted, type)) == NODQUOT)
      if ((lptr = add_dquot(wanted, type)) == NODQUOT)
	 abort_now("add_to_quota", "Can't add dquot structure type %s for uid %d\n",
		   quotatypes[type], wanted);

   /*
    * A dir is a special case, we count it for 1 inode and 1 block.
    */
   if (S_ISDIR(st->st_mode)) {
      lptr->dq_curinodes++;
      lptr->dq_curblocks++;
   } else {
      /*
       * A symlink is always counted as 1 inode and 1 block even if the 
       * filesystem is smart and stuffs it into the inode. Be fair to other
       * users that have there files on other filesystems that aren't that smart
       * or have symlinks that don't go into the inode (e.g. to long).
       */
      if (S_ISLNK(st->st_mode)) {
         lptr->dq_curinodes++;
         lptr->dq_curblocks++;
      } else {
         if (st->st_nlink != 1) {
            store_dlinks(st, type);
            return;
         }
         lptr->dq_curinodes++;
         if (st->st_size)
            lptr->dq_curblocks += isize_to_blocks(st->st_size, st->st_blksize);
      }
   }
}

void abort_now(const char *perror_mes, const char *fmt, ...)
{
   va_list args;

   va_start(args, fmt);
   vfprintf(stderr, fmt, args);
   va_end(args);

   fflush(stderr);
   perror(perror_mes);
   exit(-1);
}

void add_dlinks(u_long * inodes, u_long * blocks, int id, int type)
{
   struct dlinks  *lptr;

   if (dflag)
      fprintf(stderr, "Adding blocks from hardlinks for %s %d\n",
              quotatypes[type], id);

   for (lptr = stored_links[type]; lptr != (struct dlinks *)NULL; lptr = lptr->next) {
      if (lptr->id == id) {
         (*inodes)++;
         if (lptr->size)
            *blocks += isize_to_blocks(lptr->size, lptr->blksize);
         files_done++;
      }
   }
}

/*
 * Clean up all list from a previous run.
 */
void remove_list()
{
   int cnt;
   struct dquot *dquot, *dquot_free;
   struct dlinks *dlink, *dlink_free;

   for (cnt = 0; cnt < MAXQUOTAS; cnt++) {
      if (dquot_list[cnt] != NODQUOT) {
         dquot = dquot_list[cnt];
         while (dquot != NODQUOT) {
            dquot_free = dquot;
            dquot = dquot->next;
#ifdef DEBUG_MALLOC
	    free_mem += sizeof(struct dquot);
#endif
            free(dquot_free);
         }
         mru_dquot[cnt] = NODQUOT;
      }
      dquot_list[cnt] = NODQUOT;
      if (stored_links[cnt] != (struct dlinks *)NULL) {
         dlink = stored_links[cnt];
         while (dlink != (struct dlinks *)NULL) {
            dlink_free = dlink;
            dlink = dlink->next;
#ifdef DEBUG_MALLOC
	    free_mem += sizeof(struct dlinks);
#endif
            free(dlink_free);
         }
      }
      stored_links[cnt] = (struct dlinks *)NULL;
   }
}

/*
 * Dump the quota info that we have in memory now to the appropriate
 * quota file. We lock it during the time we update it.
 */
void dump_to_file(char *fsname, char *quotafile, int type)
{
   struct dqblk dq_dqb;
   struct dquot *dquot;
   int quota_enabled = 0, max_id;
   int fd, id = 0;

   if (vflag || dflag)
      fprintf(stderr, "Using quotafile %s\n", quotafile);

   if (quotactl(QCMD(Q_GETQUOTA, type), fsname, 0, (caddr_t)&dq_dqb) == 0)
      quota_enabled = 1;

   if ((vflag || dflag) && quota_enabled)
      fprintf(stderr, "Updating in-core %s quotas\n", quotatypes[type]);

   if ((fd = open(quotafile, O_RDWR | O_CREAT, 0600)) < 0)
      abort_now("open", "dump_to_file(%s): ", quotafile);

   if (flock(fd, LOCK_EX) < 0)
      abort_now("flock", "dump_to_file(%s): ", quotafile);

   /*
    * First dump the gracetimes that are always a first in the 
    * quotafile. Only dump new gracetimes when creating a new 
    * quotafile.
    */
   
   if (read(fd, &dq_dqb, sizeof(struct dqblk)) <= 0) {
      memset((caddr_t *)&dq_dqb, 0, sizeof(struct dqblk));
      dq_dqb.dqb_btime = MAX_DQ_TIME;
      dq_dqb.dqb_itime = MAX_IQ_TIME;
      write(fd, &dq_dqb, sizeof(struct dqblk));
   }

   max_id = highestid[type];
   while (id <= max_id) {
      if (lseek(fd, dqoff(id), SEEK_SET))
         read(fd, &dq_dqb, sizeof(struct dqblk));
      if ((dquot = lookup_dquot(id, type)) != NODQUOT) {
         dq_curinodes = dquot->dq_curinodes;
         dq_curblocks = dquot->dq_curblocks;
         if (dflag)
            fprintf(stderr, "%s %d: curinodes: %d curblocks: %d\n",
                    quotatypes[type], id, dq_curinodes, dq_curblocks);
         add_dlinks(&dq_curinodes, &dq_curblocks, id, type);
         if (dflag)
            fprintf(stderr, "%s %d: curinodes: %d curblocks: %d\n",
                    quotatypes[type], id, dq_curinodes, dq_curblocks);
      } else
         memset((caddr_t *)&dq_dqb, 0, sizeof(struct dqblk));

      /*
       * If the quota is updated with the systemcall it isn't needed to update
       * it in the file. Because the kernel will do that with the next sync.
       */
      if (quota_enabled)
         if (quotactl(QCMD(Q_SETUSE, type), fsname, id, (caddr_t)&dq_dqb) == 0) {
            id++;
            continue;
         }

      if (lseek(fd, dqoff(id), SEEK_SET))
         write(fd, &dq_dqb, sizeof(struct dqblk));
      id++;
   }
   flock(fd, LOCK_UN);
   close(fd);
}