/********************************************************************\
*								     *
*			Dave's DU version 1.01                       *
*								     *
* ©1991 by Dave Schreiber.  All rights reserved.  This program may   *
* be freely distributed.  The only charges allowed in relation to    *
* distributing this program is in relation to media, shipping, and   *
* handling costs.						     *
*								     *
* Version list: 						     *
*   1.02 (05/24/91) - Improved CTRL-C handling                       *
*   1.00 (04/29/91) - Original release version                       *
*								     *
* Compiled with SAS/C V5.10.  To compile, do the following:	     *
*   1> lc DU							     *
*   1> blink with DU.lnk					     *
*								     *
\********************************************************************/


#include <exec/types.h>
#include <exec/exec.h>
#include <exec/tasks.h>
#include <dos/dos.h>
#include <dos/dosextens.h>
#include <stdio.h>
#include <signal.h>

ULONG *DOSBase;

/*Function prototypes*/

ULONG __regargs getDirSize(BPTR lock,char *dirName);
void interpretArgs(int argc,char *argv[],WORD *countDirs,WORD *countFiles,
		    char *dirName);
void printSize(char *dirName,ULONG size);
void printHelp(void);
void ctrlC_Handler(void);

/*-------------------*/

/*Pragmas*/

#pragma syscall AllocMem C6 1002
#pragma syscall FreeMem D2 0902
#pragma syscall FindTask 126 901

#pragma libcall DOSBase CurrentDir 7E 101
#pragma libcall DOSBase Examine 66 2102
#pragma libcall DOSBase ExNext 6C 2102
#pragma libcall DOSBase Lock 54 2102
#pragma libcall DOSBase UnLock 5A 101

/*-------*/

/*Global variables*/

WORD countFiles,countDirs;
struct Task *task;
BYTE quitNow=FALSE;

/*----------------*/

/*Defines*/

#define ERROR 0xFFFFFFFF
#define ABS(x) ((x) < 0 ? -(x) : (x))

/*-------*/

/*Main.  Initializes, find the directory size, and then cleans up*/
main(int argc,char *argv[])
{
   BPTR lock;
   char dirName[256];
   LONG totalSize;

      /*Interpret the command line arguments*/
   interpretArgs(argc,argv,&countDirs,&countFiles,dirName);

      /*Replace the CTRL-C handler*/
   signal(SIGINT,ctrlC_Handler);

      /*Get a lock on the provided directory name*/
   lock=Lock(dirName,SHARED_LOCK);
   if(lock == NULL)  /*Abort if the directory isn't there*/
   {
      fputs("Couldn't access the directory you requested;  aborting.\n",
	    stderr);
      exit(200);
   }

      /*Get the size of the contents of the directory*/
   totalSize=getDirSize(lock,dirName);

      /*Abort if something went wrong, otherwise print the directory size*/
   if(totalSize == ERROR)
      fputs("Error.  Aborting!\n",stderr);
   else
      if(!countDirs && !quitNow)
	 printSize(dirName,totalSize);

   UnLock(lock);  /*Free the lock on the top directory*/
   exit(0);       /*And exit*/
}

/*Get the size of the contents of the specified directory (recursive)*/
ULONG __regargs getDirSize(BPTR lock,char *dirName)
{
   register ULONG totalSize=0;
   register struct FileInfoBlock *fib;
   register BPTR dirLock,oldCurDir;
   register UWORD dirNameLen;
   register ULONG size;

   dirNameLen=strlen(dirName);

      /*Allocate (longword aligned) FileInfoBlock structure*/
   fib=(struct FileInfoBlock *)AllocMem(sizeof(struct FileInfoBlock),
	       MEMF_PUBLIC,MEMF_CLEAR);
   if(fib==NULL)
      return(ERROR);

      /*Save the current directory, and change to the directory we need*/
      /*to find the size of.*/
   oldCurDir=CurrentDir(lock);
   Examine(lock,fib);

      /*For each file/directory, get the size...*/
   while(ExNext(lock,fib) && !quitNow)
   {
      totalSize++;	/*Add 1 (for the directory entry)*/
      if(ABS(fib->fib_EntryType) != 4)  /*If the file isn't a link...*/
      {
	 if(fib->fib_EntryType < 0)  /*If the file isn't a directory*/
	 {				     /*get the size directly*/
	       /*Add the size to the total*/
	    totalSize+=fib->fib_NumBlocks;

	    if(countFiles) /*If the user wants the size of each file*/
	    {		   /*printed, print the file size*/
	       if(dirNameLen!=0 && dirName[dirNameLen-1]!=':')
		  strcat(dirName,"/");
	       strcat(dirName,fib->fib_FileName);
	       printSize(dirName,totalSize);
	       dirName[dirNameLen]=NULL;
	    }
	 }
	 else  /*It's a directory*/
	 {
	       /*Get the lock on the directory*/
	    dirLock = Lock(fib->fib_FileName,SHARED_LOCK);

	    if(dirLock == NULL)  /*Abort if there was a problem getting*/
	    {			 /*a lock on the directory*/
	       CurrentDir(oldCurDir);
	       fputs("Couldn't lock directory:",stderr);
	       fputs(fib->fib_FileName,stderr);
	       fputs("\n",stderr);
	       FreeMem(fib,sizeof(struct FileInfoBlock));
	       return(ERROR);
	    }
	       /*Create the full name of the directory we're going to */
	       /*descend into*/
	    if(dirNameLen!=0 && dirName[dirNameLen-1]!=':')
	       strcat(dirName,"/");
	    strcat(dirName,fib->fib_FileName);

	       /*Recursively get the directory size*/
	    size=getDirSize(dirLock,dirName)-1;
	    UnLock(dirLock);

	    if(size==ERROR)   /*Abort if there was an error*/
	    {
	       CurrentDir(oldCurDir);
	       FreeMem(fib,sizeof(struct FileInfoBlock));
	       return(ERROR);
	    }
	    else
	       totalSize+=size+2;  /*Otherwise, add the size to the total*/

	       /*Restore the old directory name*/
	    dirName[dirNameLen]=NULL;
	 }
      }
   }
   CurrentDir(oldCurDir);

   if(quitNow)    /*If the user pressed CTRL-C, abort*/
      fputs("*** Break: du\n",stderr);
   else
      if(countDirs)  /*If CTRL-C wasn't pressed, print the size of the*/
	 printSize(dirName,totalSize); /*directory (if the user wants it)*/

      /*Deallocate the FileInfoBlock*/
   FreeMem(fib,sizeof(struct FileInfoBlock));
      /*Return the total size*/
   return(totalSize);
}

/*Interpret the command line arguments*/
void interpretArgs(int argc,char *argv[],WORD *countDirs,WORD *countFiles,
		    char *dirName)
{
   int i;

   if(argv[1][0]=='?')  /*The user is asking for help*/
   {
      printHelp();
      exit(0);
   }

      /*Set these variables to their default value*/
   *countDirs=TRUE;
   *countFiles=FALSE;
   strcpy(dirName,"");

   for(i=1;i<argc;i++)
   {
      if(argv[i][0]=='-')
	 switch(argv[i][1])
	 {
	    case 's':   /*Don't print the size of each directory*/
	    case 'S':
	       *countDirs=FALSE;
	       break;
	    case 'a':   /*Print the size of each directory*/
	    case 'A':
	       *countFiles=TRUE;
	       break;
	 }
      else  /*If the first character wasn't '-', the argument must be the*/
	 strcpy(dirName,argv[i]);   /*directory name*/
   }
      /*If the user specified that directory sizes wouldn't be printed,*/
      /*don't print out the file sizes either*/
   if(*countDirs==FALSE)
      *countFiles=FALSE;
   return;
}

/*Make the directory/file name, then print out the size*/
void printSize(char *dirName,ULONG size)
{
   char sizeStr[16];

      /*Conver the size (int) to a string*/
   stci_d(sizeStr,size);

      /*Print it out*/
   fputs(sizeStr,stdout);

      /*Print out a TAB*/
   fputs("\t",stdout);

      /*Print out the directory name, or '.' if it's the current directory*/
   if(dirName[0]==NULL)
      puts(".");
   else
      puts(dirName);

      /*Flush the output stream*/
   flushall();
   return;
}

/*Print help*/
void printHelp(void)
{
   puts("\nDave's DU V1.01 ©1991 by Dave Schreiber. All rights reserved.");
   puts("\nUsage:");
   puts(" du [-s][-a][<dir>]");
   puts("-s -- Print only the grand total size");
   puts("-a -- Print the size for each file");
   puts("<dir> -- Starting directory");
   puts("");
   return;
}

/*CTRL-C handler*/
void ctrlC_Handler(void)
{
   signal(SIGINT,ctrlC_Handler);
   quitNow=TRUE;
   return;
}