#!/usr/bin/perl
# cpanel - update_db_cache                        Copyright(c) 2010 cPanel, Inc.
#                                                           All rights Reserved.
# copyright@cpanel.net                                         http://cpanel.net
# This code is subject to the cPanel license. Unauthorized copying is prohibited

BEGIN { unshift @INC, '/usr/local/cpanel'; }

use strict;
use warnings;
use IO::Handle                   ();
use Cpanel::Config::Users        ();
use Cpanel::Config::FlushConfig  ();
use Cpanel::Config::LoadCpConf   ();
use Cpanel::DB::Map::Collection::Index  ();
use Cpanel::PwCache              ();
use Cpanel::MysqlUtils           ();
use Cpanel::PostgresAdmin::Check ();
use Cpanel::SafeDir::MK          ();
use Cpanel::Hostname             ();
use Cpanel::Validate::DB         ();
use DBI                          ();
use DBD::mysql                   ();
use Digest::MD5                  (); #for caching

my $cpanel_conf    = Cpanel::Config::LoadCpConf::loadcpconf();
my $NEEDDISKUSED   = exists $cpanel_conf->{'disk_usage_include_sqldbs'} ? $cpanel_conf->{'disk_usage_include_sqldbs'} : 1;
my $USEINFOSCHEMA  = exists $cpanel_conf->{'use_information_schema'} ? $cpanel_conf->{'use_information_schema'} : 1;
my $datastore_path = '/var/cpanel/datastore';

if ( !$NEEDDISKUSED ) {
    clear_db_caches();
    exit;
}

my %DBS;
my $mysqlversion = Cpanel::MysqlUtils::mysqlversion();
my $mysql_dbindex = Cpanel::DB::Map::Collection::Index->new( { 'db' => 'MYSQL' } );

my $hostname   = Cpanel::Hostname::gethostname();
my $database   = 'mysql';
my $mysqlpass  = Cpanel::MysqlUtils::getmydbpass('root');
my $mysqlhost  = Cpanel::MysqlUtils::getmydbhost('root') || 'localhost';
my $grant_host = 'localhost';
if ( $mysqlhost ne 'localhost' ) {
    $grant_host = $hostname;
}

if ( $mysqlversion >= 5 && $USEINFOSCHEMA) {

    my $mysql_dbh = DBI->connect( "DBI:mysql:mysql:$mysqlhost", 'root', $mysqlpass );
    my $query = $mysql_dbh->prepare("SELECT TABLE_SCHEMA as DB,SUM(DATA_LENGTH)+SUM(INDEX_LENGTH) AS SPACEUSED from information_schema.tables GROUP BY TABLE_SCHEMA;");
    $query->execute();
    my ( $map, $cpuser );
    while ( my $data = $query->fetchrow_hashref() ) {
        $cpuser = $mysql_dbindex->get_dbuser_by_db( $data->{'DB'} ) || 'root';
        $DBS{'mysql'}{$cpuser}{ $data->{'DB'} } = $data->{'SPACEUSED'};
    }
    $query->finish();
}
else {

    my $is_remote = Cpanel::MysqlUtils::is_remote_mysql();
    exit if $is_remote;

    my $mysql_datadir = Cpanel::MysqlUtils::getmysqldir();

    opendir(my $datadir_fh, $mysql_datadir) || do {
        print "Unable to read from $mysql_datadir\n";
        exit 1;
    };

    my @db_list = map  { [$mysql_datadir . $_, $_] }
                  grep { -d $mysql_datadir . $_  }
                  grep { !/^\.\.?/} readdir($datadir_fh);
    closedir($datadir_fh);

    foreach my $dir (@db_list) {
        my $dbname = $dir->[1];
        my $cpuser = $mysql_dbindex->get_dbuser_by_db($dbname) || 'root';
        $DBS{'mysql'}{$cpuser}{$dbname} += ( stat($dir->[0]) )[12] * 512;
        opendir( my $db_dir, $dir->[0] ) || do {
            print "Unable to read from $dir->[0]\n";
            exit 1;
        };

        while ( my $file = readdir($db_dir) ) {
            next if ( $file =~ /^\.+$/o );
            $DBS{'mysql'}{$cpuser}{$dbname} += ( stat( $dir->[0] . '/' . $file ) )[12] * 512;
        }
        closedir($db_dir);
    }
}

my $pg_check = Cpanel::PostgresAdmin::Check::is_configured();
if ( $pg_check->{'status'} ) {
    my $pg_dbindex = Cpanel::DB::Map::Collection::Index->new( { 'db' => 'PGSQL' } );
    my $rdr = IO::Handle->new();
    if ( my $pid = open( $rdr, '-|' ) ) {
        my ( $db, $diskusage, $map, $cpuser );
        while (<$rdr>) {
            chomp();
            ( $db, $diskusage ) = split( /=/, $_, 2 );
            $cpuser = $pg_dbindex->get_dbuser_by_db($db) || 'root';
            $DBS{'postgresql'}{$cpuser}{$db} = int $diskusage;
        }
    }
    else {
        my $postgres_tool = -x '/usr/local/cpanel/bin/postgrestool' ? '/usr/local/cpanel/bin/postgrestool' : '/usr/local/cpanel/bin/postgrestool.pl';
        exec $postgres_tool, 'LISTDBSWITHSPACE';
        exit;
    }
}

if ( !scalar keys %DBS ) {
    clear_db_caches();
    exit;
}


if ( !-e $datastore_path ) {
    Cpanel::SafeDir::MK::safemkdir( $datastore_path, '0755' );
}

my ( %MYSQL_DISK_USAGE, %POSTGRES_DISK_USAGE );

my %USERS       = cpusers();
my $pwcache_ref = pwcache();

foreach my $pwref (@$pwcache_ref) {
    my ( $user, $gid ) = (@$pwref)[ 0, 3 ];
    next if ( !exists $USERS{$user} );

    my $user_datastore_path = $datastore_path . '/' . $user;

    # Note: Variable file handles for MYSQL_FH and POSTGRESQL_FH not used to avoid IO::Handle::DESTROY overhead
    my $mysql_db_count = exists $DBS{'mysql'}{$user} ? scalar keys %{ $DBS{'mysql'}{$user} } : 0;
    if ($mysql_db_count) {

        if ( !-e $datastore_path . '/' . $user ) {
            mkdir $user_datastore_path, 0750;
            chown 0, $gid, $user_datastore_path;
        }

        if ( open( MYSQL_FH, '>', $user_datastore_path . '/mysql-db-count' ) ) {
            print MYSQL_FH $mysql_db_count;
            close(MYSQL_FH);
        }
        if ( open( MYSQL_DISK_FH, '>', $user_datastore_path . '/mysql-disk-usage' ) ) {
            my $bytes_total = 0;
            foreach my $db ( keys %{ $DBS{'mysql'}{$user} } ) {
                next if ( !$DBS{'mysql'}{$user}{$db} || $DBS{'mysql'}{$user}{$db} =~ /^[^0-9]+$/ );
                $bytes_total += int $DBS{'mysql'}{$user}{$db};
            }
            print MYSQL_DISK_FH int $bytes_total;
            $MYSQL_DISK_USAGE{$user} = int $bytes_total;
            close(MYSQL_DISK_FH);
        }

        if ( open( MYSQL_DB_FH, '>', $user_datastore_path . '/mysql-db-usage' ) ) {
            my $dbs_ref = $mysql_dbindex->find_dbs_by_dbuser($user);
            foreach my $db (@$dbs_ref) {
                my $usage = ( exists $DBS{'mysql'}{$user}{$db} && defined $DBS{'mysql'}{$user}{$db} ) ? $DBS{'mysql'}{$user}{$db} : 0;
                printf MYSQL_DB_FH "%s: %s\n", $db, $usage;
            }
            close(MYSQL_DB_FH);
        }
    }
    my $postgres_db_count = exists $DBS{'postgresql'}{$user} ? scalar keys %{ $DBS{'postgresql'}{$user} } : 0;
    if ($postgres_db_count) {

        if ( !-e $datastore_path . '/' . $user ) {
            mkdir $user_datastore_path, 0750;
            chown 0, $gid, $user_datastore_path;
        }

        if ( open( POSTGRESQL_FH, '>', $user_datastore_path . '/postgres-db-count' ) ) {
            print POSTGRESQL_FH $postgres_db_count;
            close(POSTGRESQL_FH);
        }
        if ( open( POSTGRES_DISK_FH, '>', $user_datastore_path . '/postgres-disk-usage' ) ) {
            my $bytes_total = 0;
            foreach my $db ( keys %{ $DBS{'postgresql'}{$user} } ) {
                $bytes_total += int $DBS{'postgresql'}{$user}{$db};
            }
            print POSTGRES_DISK_FH int $bytes_total;
            $POSTGRES_DISK_USAGE{$user} = int $bytes_total;
            close(POSTGRES_DISK_FH);
        }
    }
}

Cpanel::Config::FlushConfig::flushConfig( '/var/cpanel/datastore/mysql-disk-usage',    \%MYSQL_DISK_USAGE,    ': ' );
Cpanel::Config::FlushConfig::flushConfig( '/var/cpanel/datastore/postgres-disk-usage', \%POSTGRES_DISK_USAGE, ': ' );

# New scope created to encapsulate variables
{
    my $pwcache_ref = [];

    sub pwcache {
        if ( ref $pwcache_ref && scalar @$pwcache_ref ) {
            return $pwcache_ref;
        }
        Cpanel::PwCache::no_uid_cache(); #uid cache only needed if we are going to make lots of getpwuid calls
        Cpanel::PwCache::init_passwdless_pwcache();
        $pwcache_ref = Cpanel::PwCache::fetch_pwcache();

        return $pwcache_ref;
    }
}

{
    my %USERS = ();

    sub cpusers {
        if (%USERS) {
            return %USERS;
        }
        %USERS = map { $_ => undef } Cpanel::Config::Users::getcpusers();

        return %USERS;
    }
}

sub clear_db_caches {
    return if !-d $datastore_path;

    unlink( $datastore_path . '/mysql-disk-usage' )    if -e $datastore_path . '/mysql-disk-usage';
    unlink( $datastore_path . '/postgres-disk-usage' ) if -e $datastore_path . '/postgres-disk-usage';

    my %USERS       = cpusers();
    my $pwcache_ref = pwcache();
    foreach my $pwref (@$pwcache_ref) {
        my ( $user, $gid ) = (@$pwref)[ 0, 3 ];
        next if ( !exists $USERS{$user} );
        my $user_datastore_path = $datastore_path . '/' . $user;
        unlink grep { -e $_ } map { $user_datastore_path . '/' . $_ } ('mysql-db-usage','mysql-disk-usage','postgres-disk-usage','mysql-db-count','postgres-db-count');     
        rmdir $user_datastore_path;    # This should be safe, rmdir will fail if anything is left in the directory
    }

    rmdir $datastore_path;             # This should be safe, rmdir will fail if anything is left in the directory
}