#!/usr/bin/perl

# cpanel - bin/repair_bwsummary                   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

use strict;
use warnings;

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

use Cpanel::BandwidthDB            ();
use Cpanel::BWFiles                ();
use Cpanel::WildcardDomain         ();
use Cpanel::Config::LoadCpUserFile ();
use Cpanel::Config::Users          ();
use Cpanel::PwCache                ();
use Cpanel::SafeRun                ();
use Cpanel::Version                ();
use File::Copy                     ();

if ( Cpanel::Version::getversiononly() !~ /^11\.25/ || Cpanel::Version::getrevision() < 43473 ) {
    print "cPanel 11.25.0 Build 43473 or later is required to run this autofixer.   Forcing this autofixer to run on an older version will have unexpected results and may destroy your bandwidth data.\n";
    exit;
}

my $rrdtool_bin = '/usr/local/cpanel/3rdparty/bin/rrdtool';
my $dir         = Cpanel::BWFiles::default_dir();
my $logfh;
my $verbose = grep { '--verbose' eq $_ } @ARGV;
my $bkupdir;
my $usable_rrdtool;
my $now;
my $version = '3.3';

main() unless caller;

sub main {
    $usable_rrdtool = -x $rrdtool_bin;
    logprint( 'Begin: ', scalar localtime );
    $now = time;
    logprint('Unable to restore main domain without rrdtool support') unless $usable_rrdtool;

    $bkupdir = "$dir/backup_$version";
    mkdir $bkupdir, 0700 unless -d $bkupdir;

    my @users = Cpanel::Config::Users::getcpusers();

    foreach my $user (@users) {
        logprint( '[', scalar(localtime), "]: Checking bandwidth for user '$user'" );

        repair_user_bandwidth($user);
    }

    logprint( 'End: ', scalar localtime );
    exit 0;
}

sub _minstr {
    my $min = shift;
    foreach my $next (@_) {
        $min = $next if $next lt $min;
    }
    return $min;
}

sub repair_user_bandwidth {
    my ($user) = @_;

    my $cpuser = Cpanel::Config::LoadCpUserFile::loadcpuserfile($user);
    my @domains = map { Cpanel::WildcardDomain::safe_domain_name($_) } $cpuser->{'DOMAIN'}, @{ $cpuser->{'DOMAINS'} || [] }, @{ $cpuser->{'DEADDOMAINS'} || [] };

    if ( 1 == @domains ) {
        logprint("User $user only has one domain, no repair needed.");
        return;
    }

    my $gid = ( Cpanel::PwCache::getpwnam($user) )[3];
    restore_main_domain( $cpuser->{'DOMAIN'}, $gid );

    my $userdb = Cpanel::BandwidthDB->new( { 'name' => $user, 'dir' => $dir, 'gid' => $gid } );
    if ( exists $userdb->{'bandwidth'}->{'daily'}->{'bw_summary_repair_2'} ) {

        # already repaired.
        logprint("$user data already repaired, skipping");
        return;
    }

    # Backup the user file before changes.
    foreach my $file ( map { "$user$_" } '', qw/.hour .5min .remainder/ ) {
        File::Copy::copy( "$dir/$file", "$bkupdir/$file.$now" ) if -e "$dir/$file";
    }

    my $first = _minstr( keys %{ $userdb->{'bandwidth'}->{'hourly'} } );
    $userdb->bw_clear();

    # Restore non-http data from rrd files
    restore_from_rrd( $userdb, $user, $first, qw/smtp ftp pop3 imap/ );

    foreach my $dom (@domains) {
        my $domdb = Cpanel::BandwidthDB->new( { 'name' => $dom, 'dir' => $dir } );

        # Remove any keys newer then first
        if ( $first && $first =~ /^(\d+)\.(\d+)\.(\d+)T/ ) {
            my $bwh = $domdb->{'bandwidth'};
            foreach my $key ( keys %{ $bwh->{'daily'} } ) {
                my $cmp_key = sprintf '%d.%02d.%02d', ( $key =~ /^(\d+)\.(\d+)\.(\d+)/ )[ 2, 0, 1 ];
                if ( $key =~ /-(?:smtp|ftp|pop3|imap)/ && $cmp_key ge $first ) {
                    $bwh->{'daily'}->{$key} = 0;
                }
            }
        }
        $userdb->bw_add($domdb);
    }

    $userdb->{'bandwidth'}->{'daily'}->{'bw_summary_repair_2'} = 0;
    $userdb->write();

    # Remove outdated cache files
    unlink Cpanel::BWFiles::daily_summary_cache_file($user), Cpanel::BWFiles::bwusage_cache_file($user);

    logprint("bandwidth for $user modfied.");
    return;
}

sub restore_main_domain {
    my ( $main, $gid ) = @_;

    my $maindb = Cpanel::BandwidthDB->new( { 'name' => $main, 'gid' => $gid, 'dir' => $dir } );
    if ( exists $maindb->{'bandwidth'}->{'daily'}->{'bw_summary_repair_2'} ) {

        # already repaired.
        logprint("$main data already repaired, skipping");
        return;
    }

    # Backup the main file before changes.
    foreach my $file ( map { "$main$_" } '', qw/.hour .5min .remainder/ ) {
        File::Copy::copy( "$dir/$file", "$bkupdir/$file.$now" ) if -e "$dir/$file";
    }

    my $first = _minstr( keys %{ $maindb->{'bandwidth'}->{'hourly'} } );
    if ( !$usable_rrdtool ) {
        logprint("No rrdtool for data extraction, clearing damaged data for $main since $first.");
        clear_damaged_data( $maindb, $first );

        # Remove outdated cache files
        unlink Cpanel::BWFiles::daily_summary_cache_file($main), Cpanel::BWFiles::bwusage_cache_file($main);
        return;
    }

    if ( !restore_from_rrd( $maindb, $main, $first, qw/http/ ) ) {
        logprint("No VALID rrdtool for data extraction, clearing damaged data for $main since $first.");
        clear_damaged_data( $maindb, $first );
    }

    # Remove outdated cache files
    unlink Cpanel::BWFiles::daily_summary_cache_file($main), Cpanel::BWFiles::bwusage_cache_file($main);

    return;
}

sub restore_from_rrd {
    my ( $db, $name, $first, @types ) = @_;

    # Get earliest date
    my $start_date;
    if ( $first && $first =~ /^(\d+)\.(\d+)\.(\d+)T/ ) {
        $start_date = "$1$2$3";
        $first =~ s/T.*//;
    }
    else {
        logprint('No hourly timestamps found, no correction needed.');
        return;
    }

    my %bw;

    # http and all for 5 minute data
    foreach my $type (@types) {

        # Do 5min summaries.
        my $sum = extract_rrdtool_data( $name, $type, 'rate', $start_date );
        next unless $sum;

        while ( my ( $stamp, $bytes ) = each %{$sum} ) {
            $bw{ stamp_5min( $stamp, $type ) } = $bytes;

            # generate new all by adding up the types.
            $bw{ stamp_5min( $stamp, 'all' ) } += $bytes;
        }
    }
    if ( keys %bw ) {
        $db->{'bandwidth'}->{'5min'} = {%bw};
    }
    else {
        logprint("No 5min bandwidth found for $name, not critical");
    }

    my ( %bw_hr, %bw_day );

    my $extract_ok = 0;

    # http and all for hourly data
    foreach my $type (@types) {

        # Do hourly and daily summaries.
        my $sum = extract_rrdtool_data( $name, $type, 'peak', $start_date );
        unless ($sum) {
            logprint("Unable to extract hourly/daily data for $name, skipping data since $first");
            next;
        }

        $extract_ok++;
        while ( my ( $stamp, $bytes ) = each %{$sum} ) {
            $bw_hr{ stamp_hour( $stamp, $type ) } = $bytes;
            $bw_day{ stamp_day( $stamp, $type ) } += $bytes;

            # generate new all by adding up the types.
            $bw_hr{ stamp_hour( $stamp, 'all' ) } += $bytes;
            $bw_day{ stamp_day( $stamp, 'all' ) } += $bytes;
        }
    }
    $db->{'bandwidth'}->{'daily'}->{'bw_summary_repair_2'} = 0;

    return 0 if !$extract_ok;

    logprint("$name has valid rrdtool data.  Restored to DB");

    if ( keys %bw_hr ) {

        # need to merge here, not replace
        my $bwh = $db->{'bandwidth'};

        # copy in earlier data
        foreach my $key ( keys %{ $bwh->{'daily'} } ) {
            my $cmp_key = sprintf '%d.%02d.%02d', ( $key =~ /^(\d+)\.(\d+)\.(\d+)/ )[ 2, 0, 1 ];
            $bw_day{$key} ||= $bwh->{'daily'}->{$key} if $key !~ /-all|-http/ || $cmp_key lt $first;
        }

        $db->{'bandwidth'}->{'hourly'} = {%bw_hr};
        $db->{'bandwidth'}->{'daily'}  = {%bw_day};
    }
    $db->write(1);

    return 1;
}

sub stamp_5min {
    my ( $stamp, $type ) = @_;
    my ( $min, $hr, $day, $mon, $year ) = ( localtime $stamp )[ 1 .. 5 ];
    $mon  += 1;
    $year += 1900;
    sprintf '%d.%02d.%02dT%02d:%02d-%s', $year, $mon, $day, $hr, $min, $type;
}

sub stamp_hour {
    my ( $stamp, $type ) = @_;
    my ( $min, $hr, $day, $mon, $year ) = ( localtime $stamp )[ 1 .. 5 ];
    $mon  += 1;
    $year += 1900;
    sprintf '%d.%02d.%02dT%02d-%s', $year, $mon, $day, $hr, $type;
}

sub stamp_day {
    my ( $stamp, $type ) = @_;
    my ( $day, $mon, $year ) = ( localtime $stamp )[ 3 .. 5 ];
    $mon  += 1;
    $year += 1900;
    "$mon.$day.$year-$type";
}

sub clear_damaged_data {
    my ( $maindb, $first ) = @_;
    my $bwh = $maindb->{'bandwidth'};
    $bwh->{'5min'}   = {};
    $bwh->{'hourly'} = {};
    foreach my $key ( keys %{ $bwh->{'daily'} } ) {
        my $cmp_key = sprintf '%d.%02d.%02d', ( $key =~ /^(\d+)\.(\d+)\.(\d+)/ )[ 2, 0, 1 ];
        if ( $key =~ /-http/ && $cmp_key ge $first ) {
            my $all_key = $key;
            $all_key =~ s/http/all/;
            $bwh->{'daily'}->{$all_key} -= $bwh->{'daily'}->{$key}
              if $bwh->{'daily'}->{$all_key} >= $bwh->{'daily'}->{$key};
            $bwh->{'daily'}->{$key} = 0;
        }
    }
    $maindb->write(1);
}

sub extract_rrdtool_data {
    my ( $main, $type, $set, $start_date ) = @_;
    my $scale = 'peak' eq $set ? 3600 : 300;

    my $file = Cpanel::BWFiles::one_rrd_file( $main, $type, $set, $dir );

    my @lines = grep { !/nan/ && !/output/ && /\S/ } Cpanel::SafeRun::saferunnoerror( $rrdtool_bin, 'fetch', $file, 'AVERAGE', "--start=$start_date" );

    if ( @lines < 1 ) {
        logprint("No rrdtool data found for $main ($type, $set)");
        return;
    }
    my %summary;
    foreach my $line (@lines) {
        next unless $line =~ /^(\d+): (\S+)/;
        $summary{$1} = sprintf '%.0f', $2 * $scale;
    }
    return \%summary;
}

sub logprint {
    unless ($logfh) {
        my $logdir = '/var/cpanel/logs';
        mkdir $logdir, 0700 unless -d $logdir;
        my $logfile = "$logdir/bwsummary-autofixer." . time;
        open $logfh, '>>', $logfile or do {

            # Can't open file
            warn "Cannot open $logfile for writing: $!\n";
            $verbose = 1;
        };
        print $logfh "---------- $0: verbose: $verbose -----------------\n" if $logfh;
    }
    print $logfh @_, "\n" if $logfh;    # Handle case where we are unable to write logfile.
    print @_, "\n" if $verbose;
}