#!/usr/bin/perl
# -------------------------------------------------------------------------
# Copyright (C) 2002 TJ Saunders <tj@castaglia.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.
#
# 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, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307, USA.
#
#  $Id: diskuse,v 1.5 2002/04/27 15:23:36 tj Exp tj $
#
# -------------------------------------------------------------------------

use strict;

use Fcntl;
use File::Basename qw(basename);
use Getopt::Long;

my $program = basename($0);
my %opts = {};

GetOptions(\%opts, 'G', 'K', 'M', 'dir-tally', 'group=s', 'help',
  'user=s', 'verbose');

my $verbose = 0;
my $user_id = -1;
my $group_id = -1;

usage() if defined($opts{'help'});

if (length(@ARGV) < 1) {
  print STDOUT "$program: wrong number of parameters\n";
  exit 0;
}

unless (defined($opts{'user'}) || defined($opts{'group'})) {
  print STDOUT "$program: missing required --group or --user option\n";
  exit 0;
}

if (defined($opts{'user'})) {
  $user_id = getpwnam($opts{'user'}) or
    die "$program: no such user: $opts{'user'}\n";
}

if (defined($opts{'group'})) {
  $group_id = getgrnam($opts{'group'}) or
    die "$program: no such group: $opts{'group'}\n";
}

$verbose = 1 if (defined($opts{'verbose'}));

my ($total_bytes, $total_files);

foreach my $dir (@ARGV) {
  my ($bytes, $files) = get_dir_disk_use($dir);

  $total_bytes += $bytes;
  $total_files += $files;
}

my $byte_units = "bytes";

if (defined($opts{'K'})) {
  $total_bytes /= 1024;
  $byte_units = "KB";

} elsif (defined($opts{'M'})) {
  $total_bytes /= (1024 * 1024);
  $byte_units = "MB";

} elsif (defined($opts{'G'})) {
  $total_bytes /= (1024 * 1024 * 1024);
  $byte_units = "GB";
}

my $file_units = "files";
$file_units = "file" if ($total_files == 1);

print STDOUT "$program: $total_bytes $byte_units in $total_files $file_units\n";

# done
exit 0;

# --------------------------------------------------------------------------
sub get_dir_disk_use {
  my ($dir) = @_;

  print STDOUT "$program: examining directory '$dir'\n" if $verbose;

  my @subdirs = ();
  my $nbytes = 0;
  my $nfiles = 0;

  unless (opendir(DIR, $dir)) {
    print STDERR "$program: unable to open directory '$dir': $!\n";
    return (0, 0);
  }

  my @files = readdir(DIR);
  close(DIR);

  foreach my $file (@files) {
    my ($mode, $uid, $gid, $size) = (lstat("$dir/$file"))[2, 4, 5, 7];

    # Is this entry a regular file, or a directory?
    unless (-f "$dir/$file" || (-l "$dir/$file" && !-f "$dir/$file") || -d "$dir/$file") {
      print STDOUT "$program: skipping '$dir/$file': not a file or directory\n" if $verbose;
      next;
    }

    # Does this file match the given user/group ownership?
    if ($user_id != -1 && $group_id != -1) {

      if ($uid == $user_id && $gid == $group_id) {

        if (defined($opts{'dir-tally'}) && $file eq ".") {
          $nbytes += $size;
          $nfiles++;
 
        } elsif (! -d "$dir/$file") {
          $nbytes += $size;
          $nfiles++;
        }

      } else {
        print STDOUT "$program: '$file' does not match UID $user_id, GID $group_id\n" if $verbose;
      }

    } elsif ($user_id != -1) {

      if ($uid == $user_id) {

        if (defined($opts{'dir-tally'}) && $file eq ".") {
          $nbytes += $size;
          $nfiles++;

        } elsif (! -d "$dir/$file") {
          $nbytes += $size;
          $nfiles++;
        }

      } else {
        print STDOUT "$program: '$file' does not match UID $user_id\n" if $verbose;
      }

    } elsif ($group_id != -1) {

      if ($gid == $group_id) {

        if (defined($opts{'dir-tally'}) && $file eq ".") {
          $nbytes += $size;
          $nfiles++;

        } elsif (! -d "$dir/$file") {
          $nbytes += $size;
          $nfiles++;
        }

      } else {
        print STDOUT "$program: '$file' does not match GID $group_id\n" if $verbose;
      }
    }

    # Is this entry a directory?  If so, add it to the subdir list.
    if (-d "$dir/$file") {

      # Skip dot directories
      push(@subdirs, "$dir/$file") unless ($file eq "." || $file eq "..");
    }
  }

  # Now, recurse through the directory's subdirectories
  foreach my $subdir (@subdirs) {
    my ($bytecount, $filecount) = get_dir_disk_use($subdir);
    $nbytes += $bytecount;
    $nfiles += $filecount;
  }

  return ($nbytes, $nfiles);
}

# --------------------------------------------------------------------------
sub usage {

  print STDOUT <<END_OF_USAGE;

usage: $program [options] dir1 dir2 ... dirN

 $program calculates the amount of disk space used underneath the given
 directory (or directories).  Either a --group or a --user option is required.
 If both are used, only files whose ownership matches both the given name and
 group will be added to the disk usage sum.

 Options:

  --G          Display the number of bytes in units of gigabytes.  The default
               number is in bytes.

  --K          Display the number of bytes in units of kilobytes.  The default
               number is in bytes.

  --M          Display the number of bytes in units of megabytes.  The default
               number is in bytes.

  --dir-tally  Count the size of directories toward the disk usage sum.  The
               default is to consider regular files only.

  --group      Count files owned by this group in the disk usage sum.

  --help       Display this message.

  --user       Count files owned by this user in the disk usage sum.

  --verbose    Display diagnostic output while $program is running.

END_OF_USAGE

  exit 0;
}

# --------------------------------------------------------------------------