#!/usr/local/bin/perl

use strict;

=head1 NAME
    
disk_usage.pl

=head1 DESCRIPTION

This script reports the disk usage on an account.  This script takes in
account Dovecot index files that does not go against a user's disk usage.

disk_usage.pl is used by dbc.pl only for legacy hosting accounts. It
should produce accurate results on legacy accounts and vps's, though
vps's don't count disk space for billing purposes. 

=cut

use Getopt::Long;
use Data::Dumper;

my $GNUDU;
if ($^O eq 'linux') {
    $GNUDU = '/usr/bin/du';
} elsif ($^O eq 'freebsd') {
    $GNUDU = '/usr/local/bin/gdu';
} else {
    die "unsupported os: $^O";
}

my $TOTALS = {};

# Run CLI portion if we're not calling it from somewhere.
run() unless caller();

=head2 run

This is the portion that is run if this file is call as a script.  This will
use the funtions below to print out diskusage in a user friendly way.

=cut

sub run {

    # We get a few different keys back from the get_counts_for_user, we want
    # to only report on some of them.
    my $report_on_keys = { 
        mail => 1, 
        home => 1, 
        www  => 1, 
        misc => 1, 
        ftp  => 1, 
    };

    # Options defaults.
    my %opt = (
        u  => undef,
        h  => 1,
    );

    GetOptions (
        \%opt,
        'u=s',
        'h!',
    );

    # if we are running as a user who is not root, default it to the current
    # running user.  Ignoring anything that might of been passed into -u.
    my $uid  = $>;
    my $user = getpwuid($uid);

    $opt{u} = $user if $uid != 0;

    if (!$opt{u}) {
        print "ERROR: Please specifiy a user with the -u switch.\n";
        exit;
    }

    print "Gathering stats for $opt{u} ...\n";

    my $data = get_counts_for_user($opt{u}, 0);

    for my $key (keys %$report_on_keys ) {

        if ($opt{h}) {
            printf "%5s: %-15s\n", $key, _human_readable($data->{$key});
        } else { 
            printf "%5s: %-15s\n", $key, "$data->{$key} KB";
        }
    } 

        
    if ($opt{h}) {
        my $size = _human_readable($TOTALS->{total});
        printf "%5s: %-15s\n", "total", $size;
    } else { 
        printf "%5s: %-15s\n", "total", "$TOTALS->{total} KB";
    }


}

=head2 get_counts_for_user

Takes the username and returns a hash of disk counts for mail, web, and ftp.

=cut

sub get_counts_for_user {
    my ($uname) = @_;

    my %counts;

    # home
    $counts{home} = _getdu(path => "/usr/home/$uname");

    # www paths should only be counted if they aren't symlinks.
    my @www_paths;

    if (-e "/usr/www/users/$uname" && !-l "/usr/www/users/$uname") {
        push @www_paths, "/usr/www/users/$uname";
    }
    if (-e '/usr/wwws' && !-l '/usr/wwws' &&
        -e "/usr/wwws/users/$uname" && !-l "/usr/wwws/users/$uname") {
        push @www_paths, "/usr/wwws/users/$uname";
    } 
    $counts{www} = _getdu(path => \@www_paths);

    $TOTALS->{total} += $counts{www};

    # ftp
    $counts{ftp}      = _getdu(path => ["/usr/ftp/pub/$uname", "/usr/public_ftp/$uname"]);
    $TOTALS->{total} += $counts{ftp};
    
    # misc
    $counts{misc}     = _getdu(path => ["/var/cron/tabs/$uname", "/var/at/jobs/$uname"]);
    $TOTALS->{total} += $counts{misc};

    # mail
    $counts{mail}     = _getdu(path => ["/usr/boxes/$uname"], ignore => '*.index*');
    # We explicitly do not count /var/mail/$uname here, on legacy accounts.
    # The reason for this is: an early version of disk_usage.pl
    # accidentally failed to count /var/mail. Fixing that bug would
    # cause a lot of unexpected disk overusage charges. It was decided
    # to continue to not count this usage. The change only affects
    # legacy customers.

    # - _getdu the .imap directory, skipping .index files; add this to mail
    #   usage.
    $counts{'mail.imap.skip'} = _getdu(path => "/usr/home/$uname/.imap", ignore => '*.index*');
    $counts{mail}            += $counts{'mail.imap.skip'};
    $TOTALS->{total}         += $counts{mail};

    # - _getdu the .imap directory, not skipping .index files; subtract
    #   this from the home usage.
    $counts{'mail.imap'} = _getdu(path => "/usr/home/$uname/.imap");
    $counts{home}       -= $counts{'mail.imap'};

    $TOTALS->{total}    += $counts{home};

    return \%counts;

}

=head2 _getdu 

_getdu takes a directory or directories and returns the total disk space
used. It takes the following named parameters:

=over

=item path

Either a single path, or an arrayref of paths, to calculate usage for.

=item ignore

Either a single glob, or an arrayref of globs, to pass to (g)du's
--exclude.

=back

When successful, returns a number representing the disk space used, in
kilobytes. On error, warns and returns 0.

=cut 

sub _getdu {
    my %args = @_;

    my @ignore;
    if ($args{ignore}) {
        my $i = ref $args{ignore} ? $args{ignore} : [$args{ignore}];
        @ignore = map {"--exclude=$_"} @$i;
    }
    $args{path}
        or die "You have to provide a path";
    my $p = ref $args{path} ? $args{path} : [$args{path}];
    # do this so that we can trust the return value of du more;
    # otherwise, if some extra path, like /usr/wwws/users/$user, doesn't
    # exist, we'd fail
    my @paths = grep {-e $_} @$p;
    @paths
        or return 0;

    my $tot = 0;

    my @du_cmd = ($GNUDU, '-Dksc', @ignore, @paths);
    my $du_cmd = join(' ', @du_cmd, '2>/dev/null');
    my $du;
    unless (open $du, '-|', $du_cmd) {
        warn "Failed to run @du_cmd: $!";
        return 0;
    }
    while (defined (my $line = <$du>)) {
        if ($line =~ /^(\d+)\ttotal$/) {
            $tot = $1;
            last;
        }
    }

    return $tot;
}

=head2 _human_readable

A helper function that prints out the disk usage in a human friendly 
format.

=cut

sub _human_readable {
    my $data = shift;

    # we already have results in K
    my $M = 1024;
    my $G = $M * 1024;

    return "0 KB" unless $data;

    # Hack.  We only show Megabytes for disk usage
    if ((my $g = $data/$G) >= 1) {
       return sprintf "%0.2f GB", $g;
    } elsif ((my $m = $data/$M) >= 1) {
       return sprintf "%0.2f MB", $m;
    } else {
        return sprintf "%0.2f KB", $data;
    }

}

1;