#!/usr/bin/perl -w
my $VERSION = "2.01";

use Cwd 'abs_path';
use Sys::Syslog qw(:standard :macros);
use Digest::MD5 qw(md5_hex);
use strict;

# Constants (may be overriden in config).
our $CONF_OVER      = '/etc/vzfailcnt.conf';
our $MAIL_TO        = undef;
our $MAIL_FROM      = undef;
our $THRESHOLD_DISK = '80%';
our $SENDMAIL_PATH  = '/usr/sbin/sendmail';
our $UBC_FILE       = '/proc/user_beancounters';
our $UBC_PREV       = '/var/spool/user_beancounters.old';
our $CONF_DIR       = '/etc/vz/conf';
our $SCRIPT_PATH    = abs_path(__FILE__);

sub main {
    $ENV{PATH} = "$ENV{PATH}:/bin:/sbin:/usr/bin:/usr/sbin:/usr/local/bin:/usr/local/sbin";

    if (-f $CONF_OVER) {
        read_conf_over($CONF_OVER);
    }

    my $mail_to = $ARGV[0];
    if (!$mail_to || $mail_to eq "conf") {
        $mail_to = $MAIL_TO;
        usage() if !$mail_to;
    }
    $mail_to =~ /^[^@\s]+@[^@\s]+$/s or die "Invalid e-mail format: $mail_to\n";
    my $mail_from = $MAIL_FROM || $mail_to;
    $THRESHOLD_DISK =~ s/\D+//sg;

    my ($cur, $cur_content) = read_ubc_file($UBC_FILE);
    my $prv = read_ubc_file($UBC_PREV, 1);
    my $diff = calc_diff($prv, $cur);
    if (keys %$diff) {
        my $mail = get_mail_for_diff($diff, $mail_to, $mail_from);
        sendmail($mail_to, $mail);
        openlog("vzfailcnt", "ndelay,pid", LOG_USER);
        syslog(LOG_INFO, "E-mail with failcnt changes is sent to %s", $mail_to);
        closelog();
    }

    # Always save the file, even if no diffs are detected, because there may
    # be other resources appeared since the previous run, e.g. new disks are
    # mounted, new machines are started etc.
    write_file($UBC_PREV, $cur_content);
}

sub read_conf_over {
    my ($f) = @_;
    open(local *F, $f) or die "Cannot open $f: $!\n";
    while (<F>) {
        s/#.*//sg;
        my ($k, $v) = /^\s*(\w+)\s*=\s*(.*?)\s*$/s or next;
        no strict 'refs';
        $$k = $v;
    }
}

sub read_ubc_file {
    my ($fn, $is_saved_file) = @_;
    my $file_content = "";
    my @col_names = ();
    my $cur_uid = undef;
    my $cur_host = undef;
    my %result = ();
    open(local *F, $fn) or return ($is_saved_file? undef : die "Cannot open $fn: $!\n");
    foreach my $line (<F>) {
        $file_content .= $line;
        $line =~ s/^\s+|\s+$//sg;
        if ($line =~ /^\D+:/) {
            next;
        }
        if ($line =~ /^uid\s+(.*)/s) {
            @col_names = $1 =~ /(\S+)/g;
            next;
        }
        my @rows = ();
        if ($line =~ /^(\d+):\s*(.*)/s) {
            ($cur_uid, $line) = ($1, $2);
            $cur_host = get_host_by_uid($cur_uid);
            if (!$is_saved_file) {
                push @rows, get_disk_limits($cur_uid);
                foreach my $row (@rows) {
                    $file_content .= (" " x 12) . join("    ", @$row{@col_names}) . "\n";
                }
            }
            # no "next" - continue processing
        }
        if ((my @data = $line =~ /(\S+)/g) > 3) {
            my %row;
            @row{@col_names} = @data;
            push @rows, \%row;
        }
        foreach my $row (@rows) {
            $result{$cur_uid} ||= {
                hostname  => $cur_host,
                uid       => $cur_uid,
                resources => {},
            };
            $result{$cur_uid}{resources}{$row->{resource}} = $row;
        }
    }
#    use Data::Dumper; die Dumper(\%result);
    return wantarray? (\%result, $file_content) : \%result;
}

sub get_disk_limits {
    my ($uid) = @_;
    return () if $THRESHOLD_DISK !~ /^\d+$/s;
    if (!$uid) {
        return get_disk_limits_host();
    } else {
        return get_disk_limits_uid($uid);
    }
}

sub get_disk_limits_uid {
    my ($uid) = @_;
    my $lines = `vzquota stat "$uid" -f`;
    my @col_names = ();
    my @rows = ();
    foreach (split /\n/, $lines) {
        s/^\s+|\s+$//sg;
        if (!@col_names) {
            @col_names = /(\S+)/g;
            next;
        }
        if ((my @data = /(\S+)/g) > 3) {
            my %row;
            @row{@col_names} = @data;
            my $ubc_name = undef;
            if ($row{resource} eq "1k-blocks") {
                $ubc_name = "diskspace.%";
            }
            if ($row{resource} eq "inodes") {
                $ubc_name = "diskinodes.%";
            }
            my $limit = $row{softlimit} < $row{hardlimit}? $row{softlimit} : $row{hardlimit};
            if ($ubc_name && $limit) {
                my $percent = int($row{usage} / $limit * 100);
                push @rows, {
                    uid      => $uid,
                    resource => $ubc_name,
                    held     => $percent . "%",
                    maxheld  => $percent . "%",
                    barrier  => $THRESHOLD_DISK . "%",
                    limit    => "100%",
                    failcnt  => $percent > $THRESHOLD_DISK? ($percent - $THRESHOLD_DISK) : 0,
                };
            }
        }
    }
    return @rows;
}

sub get_disk_limits_host {
    my @rows = ();
    my %commands = (
       "diskspace.%"  => "df -P",
       "diskinodes.%" => "df -iP",
    );
    while (my ($ubc_name, $cmd) = each %commands) {
        my $lines = `$cmd | tail -n+2`;
        foreach (split /\n/, $lines) {
            s/^\s+|\s+$//sg;
            next if ! /^\s*(\S+)\s+(\d+)\s+(\d+)/s;
            my ($dev, $limit, $held) = ($1, $2, $3);
            next if !$limit;
            my $percent = int($held / $limit * 100);
            push @rows, {
                uid      => 0,
                resource => $ubc_name . "." . $dev,
                held     => $percent . "%",
                maxheld  => $percent . "%",
                barrier  => $THRESHOLD_DISK . "%",
                limit    => "100%",
                failcnt  => $percent > $THRESHOLD_DISK? ($percent - $THRESHOLD_DISK) : 0,
            };
        }
    }
    return @rows;
}


sub calc_diff {
    my ($prv, $cur) = @_;
    my %diff = ();
    foreach my $uid (sort keys %$cur) {
        my %res_prv = %{$prv->{$uid}{resources} || {}};
        my %res_cur = %{$cur->{$uid}{resources}};
        foreach my $res_name (sort keys %res_cur) {
            my $sub = $res_cur{$res_name}{failcnt} - ($res_prv{$res_name}{failcnt} || 0);
            next if $sub <= 0;
            $diff{$uid} ||= {
                hostname  => $cur->{$uid}{hostname},
                uid       => $uid,
                resources => {},
            };
            $diff{$uid}{resources}{$res_name} = { %{$res_cur{$res_name}}, failcntdiff => $sub };
        }
    }
#    use Data::Dumper; die Dumper(\%diff);
    return \%diff;
}

sub get_mail_for_diff {
    my ($diff, $arg_mail_to, $arg_mail_from) = @_;
    my @cols = ('uid', 'hostname', 'resource', 'failcntdiff', 'failcnt', 'held', 'maxheld', 'barrier', 'limit');
    my @html = "<table>";
    push @html, "<tr>";
    foreach my $col (@cols) {
        push @html, "<td><b>$col</b>&#160;</td>";
    }
    push @html, "</tr>";
    my (@uid_res_names, %res_names);
    foreach my $uid (sort keys %$diff) {
        foreach my $res_name (sort keys %{$diff->{$uid}{resources}}) {
            $res_names{$res_name}++;
            push @uid_res_names, "$uid:$res_name";
            push @html, "<tr>";
            my $seen_failcnt = 0;
            foreach my $col (@cols) {
                my $val = $diff->{$uid}{$col} || $diff->{$uid}{resources}{$res_name}{$col};
                if ($col eq "failcntdiff") {
                    $val = "<font color=\"#A00\"><b>+$val</b></font>";
                    $seen_failcnt = 1;
                } elsif ($seen_failcnt) {
                    $val = "<font color=\"#CCC\">$val</font>";
                }
                push @html, "<td>$val&#160;</td>";
            }
            push @html, "</tr>";
        }
    }
    push @html, '</table>';
    my $tmpl = join("", <DATA>);
    # Expand the template.
    my $mail_from = $arg_mail_from;
    my $mail_to   = $arg_mail_to;
    my $mail_html = join("\n", @html, "");
    my $myhostname = `hostname`; $myhostname =~ s/\s+//sg;
    my $res_names = join(", ", sort keys %res_names);
    my $msgid = md5_hex("$myhostname." . join(".", @uid_res_names)) . "\@vzfailcnt";
    $tmpl = eval("return qq\x01$tmpl\x01");
    #die $tmpl;
    return $tmpl;
}

sub sendmail {
    my ($to, $body) = @_;
    open(P, "|-", $SENDMAIL_PATH, "-t", $to) or die "Cannot execute $SENDMAIL_PATH: $!\n";
    print P $body;
    close(P);
}

sub get_host_by_uid {
    my ($uid) = @_;
    if (!$uid) {
       my $hostname = `hostname`;
       $hostname =~ s/\s+//sg;
       return $hostname;
    }
    if (open(local *F, "$CONF_DIR/$uid.conf")) {
        while (<F>) {
            /^\s*HOSTNAME\s*=\s*[\"\']?([^\"\'\s]+)/m and return $1;
        }
    }
    return "[unknown]";
}

sub usage {
    die join "\n", (
        "dkLab vzfailcnt v$VERSION: send OpenVZ failcnt changes over e-mail",
        "Homepage: http://en.dklab.ru/vzfailcnt",
        "Configuration: $CONF_OVER (optional)",
        "Usage:",
        "  $SCRIPT_PATH your\@email.com",
        "  - or -",
        "  $SCRIPT_PATH    # if you've defined MAIL_TO in $CONF_OVER",
        "",
        "Add script to your crontab using:",
        "  echo '*/1 * * * * root $SCRIPT_PATH' > /etc/cron.d/vzfailcntcron",
        ""
    );
}

sub write_file {
    my ($to, $content) = @_;
    open(local *T, ">", $to) or die "Cannot open $to for writing: $!\n";
    print T $content;
    close T;
}

main();

__DATA__
From: "dkLab vzfailcnt" <$mail_from>
To: <$mail_to>
Reply-To: <$mail_to>
Subject: OpenVZ failcnt increased at $myhostname: $res_names
Message-Id: <$msgid>
In-Reply-To: <$msgid>
Content-Type: text/html; charset=utf-8

<p>Hello.</p>

<p>Fail counter of resource(s) has been increased.</p>

<p>$mail_html</p>

<p>-- <br/>dkLab vzfailcnt v$VERSION</p>
