#!/usr/local/bin/perl -w
use strict;

## Scott Wiersdorf
## Created: Tue Jul 31 16:22:26 MDT 2001
## $Id: qu,v 1.9 2001/08/27 22:00:52 scottw Exp $

use vars qw($VERSION);
$VERSION = '1.06';

## qu
## quota usage

## -- use some useful modules -- ##
use Getopt::Long;
use File::Find;
use Cwd qw(abs_path getcwd);
use constant BLK_SIZE => 512;   ## real disk blocksize (I think this is constant)

## -- command-line processing -- ##
my %opt = (
	   'blocksize'    => 1024,
	   'debug'        => '',
	   'home'         => ( (getpwuid($<))[2]   ## we do special things for uid == 0
			       ? (getpwuid($<))[7] || $ENV{'HOME'}
			       : '/' ),
	   'levels'       => 2,
	   'sort'         => 'alpha',
	   'user'         => (getpwuid($<))[0] || $ENV{'LOGNAME'},
	  );
GetOptions( \%opt,
	    'help|h',
	    'verbose|v',
	    'debug:s',
	    'blocksize=i',
	    'bytes',
	    'cumulative',
	    'home=s',
	    'fast',
	    'levels:i',
	    'sort=s',
	    'summary',
	    'twiddle',
	    'user=s',
	    'version',
	  ) or usage( 'Invalid option' );

usage() if $opt{'help'};
if( $opt{'version'} ) {
    print "This is qu version $VERSION\n";
    exit;
}

## -- internal state defaults -- ##
my $home     = abs_path($opt{'home'});
my $uid      = (getpwnam($opt{'user'}))[2];
usage("user '" . $opt{'user'} . "' not found!\n")
  unless defined($uid);

my $bytes    = 0;
my $t_bytes  = 0;
my $blocks   = 0;
my $t_blocks = 0;
my %links    = ();
my %subdir   = ();
my $path     = undef;
my $path_dir = undef;
my $cwd      = getcwd();
my @args     = ();
my $cur_path_len = 0;

## leaves only $opt{'levels'} levels of directories in the path
## - if $opt{'levels'} == 1, then we print nothing beyond the path
## - if $opt{'levels'} == 0, then we do not use the $strip_level and
##   print all directories below $path
my $strip_level    = ( $opt{'levels'}
		       ? $opt{'levels'} -1
		       : 0 );
my $strip_level_re = qr!(/[^/]+(/[^/]+){0,$strip_level})!;

## -- pre-process path arguments -- ##
if( scalar(@ARGV) ) {
    @args = map { ($_ !~ m!^/! ? $cwd . '/' : $home) . $_ } @ARGV;
    for my $arg ( @args ) {
	$arg =~ s!^//!/!g;
    }
}

## use $cwd if no arguments are specified
else {
    push @args, $cwd;
}

if( $opt{'debug'} ) {
    print "\$cwd      => $cwd\n" if $opt{'debug'};
    print "\@args     => @args\n" if $opt{'debug'};
    print "\$home     => '$home'\n";
    print "\$uid      => '$uid'\n";
    print "\$level    => '" . ($opt{'summary'} ? 'undef' : $opt{'levels'}) . "'\n";
}

## -- calculate total disk usage -- ##
for my $path ( @args ) {
    $bytes   = $blocks = 0;                 ## reset totals
    %links   = %subdir  = ();               ## reset internal tracking hashes

    ## assign paths according to rules above
    $path_dir = $path;
    $path_dir =~ s!^(.*)/[^/]*$!$1! unless -d $path_dir;

    ## -- some debugging -- ##
    if( $opt{'debug'} ) {
	print "\$path     => '$path'\n";
	print "\$path_dir => '$path_dir'\n";
    }

    ## -- calculate disk usage -- ##
    printf( "%12s   %12s   %12s   %s\n", "Total Bytes", "Real size", "Bytes", "Path" ) 
      if $opt{'verbose'};
    find( \&tract, $path );                 ## this is where the action is

    ## make canonical path for non-root directories
    if( -d $path && length($path) ) {
	$path = abs_path($path);
	$path .= '/';
    }

    ## -- twiddle off -- ##
    if( $opt{'twiddle'} && !$opt{'verbose'} ) {
	print "\r" . ' 'x(24+$cur_path_len);
    }

    ## -- report disk usage -- ##
    print "\nDisk Usage for $path\n";
    unless( $opt{'summary'} ) {
	printf( "%20s   %s\n", "Disk blocks", "Directory Path" );
	printf( "%20s   %s\n", "===========", "==============" );
	for my $dir ( sort sort_order keys %subdir ) {
	    printf "%20u   %s\n", int(($subdir{$dir}*BLK_SIZE)/$opt{'blocksize'}), $dir;
	}
    }
    printf( "Blocks: %12u\n", (($blocks*BLK_SIZE)/$opt{'blocksize'}) );
    printf( "Bytes:  %12s\n", $bytes ) if $opt{'bytes'};
}

print "\nTotal Blocks: " . (($t_blocks*BLK_SIZE)/$opt{'blocksize'}) . "\n";

exit;

## hehe...
sub tract {
    my (undef,$_inode,undef,$_nlink,$_uid,undef,undef,
	$_size,undef,undef,undef,undef,$_blocks) = lstat($_);
    $_size = sprintf("%u", -s _);

    ## skip files that aren't ours
    return 0 unless $_uid == $uid;

    ## skip all but the first hard link found; we believe directories
    ## always have a link count >1 too.
    if( $_nlink > 1 ) {
	return 0 if $links{$_inode};
	$links{$_inode}++;
    }

    my $cur_path = $File::Find::dir;        ## use cwd
    $cur_path =~ s!^$path_dir!!;            ## strip start directory

    ## debug
    print "\$path_dir => $path_dir\n"
      if $opt{'debug'} eq 'cumulative';
    print "\$cur_path => $cur_path\n"
      if $opt{'debug'} eq 'cumulative';

    ## preserve only $opt{'levels'} directories. If no level is given
    ## (i.e., 0), we print all directories below $path
    if( $opt{'levels'} && !$opt{'summary'} ) {
	$cur_path =~ m/$strip_level_re/o;
	$cur_path = ( $1 ? $1 : '' );
    }
    $cur_path .= '/';                       ## add trailing slash

    ## debug
    printf( "\n(%-70s  %-16s) %u\n", $File::Find::name, $cur_path, $_blocks )
      if $opt{'debug'} eq 'cumulative';

    ## cumulative directory totals
    if( $opt{'cumulative'} && !$opt{'summary'} ) {
	## update this directory and all directories above it
	my @dirs = split( /\//, $cur_path );

	## the order of these next two lines is significant
	@dirs = map { $_ . '/' } @dirs;     ## add some trailing
					    ## slashes so we're
	                                    ## matching the same
					    ## directory path that the
                                            ## 'update totals' block does
	pop @dirs;                          ## pop the top (the top
                                            ## directory gets done after this block)

	print "scalar dirs = '" . scalar(@dirs) . "'; \$cur_path = '$cur_path'; DIRS: '@dirs'\n"
	  if $opt{'debug'} eq 'cumulative';

	my $t_path    = '';
	while( defined(my $t_segment = shift @dirs) ) {
	    $t_path .= $t_segment;
	    $subdir{$t_path} += $_blocks;
	    print "Adding $_blocks to $t_path (sum: $subdir{$t_path})\n"
	      if $opt{'debug'} eq 'cumulative';
	}
    }

    ## update totals
    $bytes   	       += $_size;           ## bytes
    $t_bytes           += $_size;           ## total bytes

    $blocks 	       += $_blocks;         ## blocks
    $t_blocks 	       += $_blocks;         ## total blocks

    ## update path statistics
    unless( $opt{'summary'} ) {
	unless( defined($subdir{$cur_path}) ) {
	    $subdir{$cur_path} = 0;
	}
	$subdir{$cur_path} += $_blocks;     ## blocks for this directory
    }

    ## twiddle our thumbs
    if( $opt{'twiddle'} && !$opt{'verbose'} ) {
	printf("\rTotal: %12u   %-${cur_path_len}s", (($t_blocks*BLK_SIZE)/$opt{'blocksize'}), $cur_path);
	$cur_path_len = length($cur_path);
    }

#    print "$bytes   $_size   $File::Find::name\n" if $opt{'verbose'};
    printf( "%12s   %12u   %12s   %s\n", $bytes, $_blocks*BLK_SIZE, $_size, $File::Find::name )
      if $opt{'verbose'};

    return 1;
}

sub sort_order {
    if( $opt{'sort'} eq 'alpha' ) {
	$a cmp $b;
    }

    elsif( $opt{'sort'} eq 'alpha_i' ) {
	lc($a) cmp lc($b);
    }

    elsif( $opt{'sort'} eq 'alpha_reverse' ) {
	$b cmp $a;
    }

    elsif( $opt{'sort'} eq 'alpha_reverse_i' ) {
	lc($b) cmp lc($a);
    }

    elsif( $opt{'sort'} eq 'alpha_breadth' ) {
	($a =~ tr!/!/!) <=> ($b =~ tr!/!/!) ||
	  $a cmp $b;
    }

    elsif( $opt{'sort'} eq 'alpha_breadth_i' ) {
	($a =~ tr!/!/!) <=> ($b =~ tr!/!/!) ||
	  lc($a) cmp lc($b);
    }

    elsif( $opt{'sort'} eq 'blocks' ) {
	$subdir{$a} <=> $subdir{$b} ||
	  $a cmp $b;
    }

    elsif( $opt{'sort'} eq 'blocks_reverse' ) {
	$subdir{$b} <=> $subdir{$a} ||
	  $a cmp $b;
    }

    else {
	$a cmp $b;
    }
}

sub usage {
    my $msg = shift;
    if( $msg ) {
	print chr(7) . "\n";
	print "ERROR: $msg";
	sleep 1;
    }

    print <<_USAGE_;

usage: qu [options] [/path1 /path2 ...]

qu calculates disk quota usage based on the specified paths (relative to
your home directory). If no path is specified, usage calculations will
be based on the current working directory.

Options:
    --help              show this menu.
    --verbose           more info than you want ;o)
    --twiddle           show block count and path while processing.
    --blocksize=n       use n when calculating block sizes. The
			default value of n is 1024 bytes so that
			qu will match the output of the
			'quota' system command. If you want 'du'-style
			512 byte block reporting, set blocksize to
			512.
    --bytes             show total bytes used in the summary report.
			Because of the layout of the UNIX filesystem,
			actual disk usage is rounded up to the nearest
			block (512 bytes).  So a file that is 999
			bytes long will actually occupy two blocks of
			512 bytes each or 1024 actual bytes.
    --cumulative        show traditional 'du'-style cumulative
                        directory usage.
    --levels=n          show n levels of directories. Not specifying
			any number (or specifying 0) will show all
			directory levels. If the 'summary' option is
			specified, no directory output will be
			displayed regardless of what 'level' is set
			to.
    --sort=sortopt      sorts the resulting directory data. If the
			'summary' option is specified, no directory
			output will be displayed regardless of what
			'sort' is set to. Valid sort options are:
			    - alpha: sort by directory alphanumerically
			    - alpha_i: like 'alpha' but case-insensitive
			    - alpha_reverse: like 'alpha' but reversed
			    - alpha_reverse_i: like 'alpha_reverse' but case-insensitive
			    - alpha_breadth: sort by "deepness" 
			    - alpha_breadth_i: like previous but case-insensitive
			    - blocks: sort by block usage
			    - blocks_reverse: like previous but reversed
    --summary           create a summary-only report; that is, do not
			show any directory information, just a terse
			disk usage statement for the specified
			directory (or directories).
    --user=user         count files owned by 'user'. By default, user
			is set to the owner of the process running qu
    --version           display version number

_USAGE_
    exit;
}

1;
__END__

=head1 NAME

B<qu> - display quota usage statistics

=head1 SYNOPSIS

B<qu> [--options] [directory ...]

=head1 DESCRIPTION

The B<qu> utility displays the file system block usage for each
directory in the file hierarchy rooted in each directory argument
owned by the current process owner (or 'user' if specified). If no
directory is specified, the block usage of the hierarchy rooted in
the I<current> directory usage is displayed. Usage is only calculated
for files and directories owned by the process owner (i.e., you),
making this useful for calculating real quota usage (hence B<qu>).

You may view statistics for multiple, disjointed file hierarchies by
specifying multiple directory arguments. You may also specify a
variety of command-line arguments to alter the output of B<qu>. See
L<"OPTIONS">.

If a path is prefixed with a leading slash ('/'), B<qu> will assume
you want to calculate disk space used relative to your home directory.
If you want to calculate usage outside of your home directory, chdir
to that directory and run B<qu>. If you type this, regardless of your
current working directory:

    qu /usr/local/lib

B<qu> will calculate disk space used in F<$HOME/usr/local/lib>.

Paths not prefixed with a leading slash will be assumed relative to
the current working directory. If your current working directory path
is F</usr/local>, typing:

    qu lib

will calculate space used by you in the real F</usr/local/lib> (i.e.,
not relative to your home directory).

B<qu> may be run by root (UID 0) with the B<--user> option to
calculate quota usage for a particular user. Unless you are root, the
B<--user> option will not likely return anything useful unless the
directories owned by the specified user are readable by your UID.

When run as root, B<qu> uses absolute filesystem paths:

    qu /

will calculate usage from the filesystem F</>, not root's home
directory (which may or may not be F</>).

=head1 OPTIONS

A variety of useful options are available for B<qu>, allowing you to
customize how B<qu> calculates its disk usage and how that information
is displayed.

Options may be specified with a single dash or the recommended
GNU-style double-dash for clarity. Options may be abbreviated to the
fewest unique letters required to avoid ambiguity:

B<qu> -bl 512 -by C<-c> /usr/local

and:

B<qu> --bl=512 --by --c /usr/local

are the same as:

B<qu> --blocksize=512 --bytes --cumulative /usr/local

Pay attention that as new options are added to B<qu>, more characters
may be required to avoid ambiguity.

=over 4

=item B<help>

Show a brief help menu for B<qu> and exit.

=item B<version>

Display the version number of B<qu> and exit.

=item B<verbose>

Show far more information than you really want ;o). This is useful if
you suspect that B<qu> has a bug and want to be able to trace it
somewhat. Seldom useful for the average user.

=item B<twiddle>

Show the current block count and path during processing.

=item B<blocksize=n>

Use n when calculating block sizes. The default value of n is 1024
bytes so that B<qu> will match the output of the 'quota' system
command. If you want 'du'-style 512 byte block reporting, set
blocksize to 512.

=item B<bytes>

Display total bytes used in the summary report. Because of the layout
of the UNIX filesystem, actual disk usage is rounded up to the nearest
block (512 bytes).  So, for example, a file that is 513 bytes long
will actually occupy two blocks of 512 bytes each or 1024 real bytes.

=item B<cumulative>

Displays the total disk usage for each directory of all blocks rooted
in that directory. This is the default 'du'-style usage.

=item B<levels=n>

Show I<n> levels of directories. Not specifying any number (or
specifying 0) will show all directory levels. If the B<summary> option
is specified, no directory output will be displayed regardless of what
B<levels> is set to.

=item B<sort=sortoption>

Sorts the resulting directory data. If the B<summary> option is
specified, no directory output will be displayed regardless of what
B<sort> is set to. Valid sort options follow.

=cut

## fix this pod section 

=pod

=over 4

=item B<alpha>

Sort by directory alphanumerically

=item B<alpha_i>

Like B<alpha> but case-insensitive

=item B<alpha_reverse>

Like B<alpha> but reversed

=item B<alpha_reverse_i>

Like B<alpha_reverse> but case-insensitive

=item B<alpha_breadth>

Sort by how "far" the directory is from the root, with the "closest"
directories first in alphanumeric order

=item B<alpha_breadth_i>

Like B<alpha_breadth> but case-insensitive

=item B<blocks>

Sort by block usage

=item B<blocks_reverse>

Like B<blocks> but reversed. This is a useful sort option ;o)

=back

=item B<summary>

Display only summary statistics; that is, do not show any directory
information, just a terse disk usage statement for the specified
directory (or directories).

=item B<user=user>

Count files owned by I<user>. By default, user is set to the owner of
the process running B<qu>.

=back

=head1 EXAMPLES

=head2 Show brief statistics for the file hierarchy rooted in the current working directory

qu --summary

=head2 Show brief statistics for my home directory (should be close to 'quota C<-v>' output)

qu --summary /

=head2 Show detailed statistics for two directory levels deep in my home directory (levels=2 is the default)

qu /

=head2 Show cumulative detailed statistics for three directory levels deep in my home directory

qu --levels=3 --cumulative /

qu C<-l> 3 C<-c> /

=head2 Show detailed statistics for ~/usr/local/lib

qu /usr/local/lib

=head2 Show detailed statistics for ~/usr/local/lib and sort the output based on which directories use the most blocks

qu --sort=blocks_reverse /usr/local/lib

=head2 Show detailed statistics for Bob's /usr/local/bin

qu --user=bob ~bob/usr/local/bin

=head1 AUTHOR

Scott Wiersdorf <scott@perlcode.org>

=head1 SEE ALSO

perl(1).

=cut