#!/usr/bin/perl -w

use strict;
use warnings;
use JSON;
use File::Which;
use Getopt::Long;

my $json = {};
my $pool = undef;
my $what = undef;
my $pretty = 0;

GetOptions(
  "zpool|pool=s" => \$pool,
  "what=s" => \$what,
  "pretty" => \$pretty
);

my $zpool = which('zpool');

if ($what and !$pool){
  print <<_EOF;
Usage:
$0 [--zpool=<name>] [--what=<item>]

<name> is an optional zpool name. If specified, will only output info for this zpool.
<item> is one of size, alloc, frag, cap, dedup, health and if specified, will only output the corresponding value
If --what is specified then --zpool is mandatory.
The default (with no option) is to output all the info of all the zpool in a JSON format
_EOF
  exit 1;
}

if ($zpool){
  my $cmd = "$zpool list -p -H" . ( ($pool) ? " $pool" : "");
  foreach (qx($cmd)){
    #NAME    SIZE  ALLOC   FREE  EXPANDSZ   FRAG    CAP  DEDUP  HEALTH  ALTROOT
    #rpool   464G  7.49G   457G         -     2%     1%  1.00x  ONLINE  -
    if (m/^(?<pool>\w+)\s+(?<size>\d+)\s+(?<alloc>\d+)\s+(?<free>\d+)\s+.+\s+(?<frag>\d+(\.\d+)?)\s+(?<cap>\d+(\.\d+)?)\s+(?<dedup>\d+(\.\d+)?)\s+(?<health>\w+)/){
      $json->{$+{pool}}->{$_} = $+{$_} foreach (grep { $_ ne 'pool' } keys %+);
    }
    $json->{$+{pool}}->{errors} = get_zpool_errors($+{pool});
    $json->{$+{pool}}->{stats}  = get_zpool_stats($+{pool});
  }
}
if ($what){
  print ((defined $json->{$pool}->{$what}) ? $json->{$pool}->{$what} : 'ZBX_UNSUPORTED');
} elsif ($pool){
  print ((defined $json->{$pool}) ? to_json($json->{$pool}, { pretty => $pretty }) : 'ZBX_UNSUPORTED');
} else {
  print to_json($json, { pretty => $pretty });
}
print "\n";
exit 0;

sub get_zpool_errors {
  my $pool = shift;
  my $errors = {
    read_errors  => 0,
    write_errors => 0,
    cksum_errors => 0
  };
  my $i = 0;
  my $index = {};
  foreach my $line (qx($zpool status $pool)){
    # Output looks like
    #   pool: rpool
    #  state: ONLINE
    # status: One or more devices has experienced an unrecoverable error.  An
    #         attempt was made to correct the error.  Applications are unaffected.
    # action: Determine if the device needs to be replaced, and clear the errors
    #         using 'zpool clear' or replace the device with 'zpool replace'.
    #    see: http://zfsonlinux.org/msg/ZFS-8000-9P
    #   scan: scrub repaired 0B in 0h5m with 0 errors on Tue May 29 10:04:31 2018
    # config:
    #
    #         NAME        STATE     READ WRITE CKSUM
    #          rpool       ONLINE       0     0     0
    #            mirror-0  ONLINE       0     0     0
    #              sda2    ONLINE       0     0     0
    #              sdb2    ONLINE       0     0   474
    #
    # errors: No known data errors

    # We want to save status, action, scan and errors
    if ($line =~ m/^\s*(scan|action|status|errors):\s+(\w+.*)/){
      $errors->{$1} = $2;
      $index->{$i} = $1;
    } elsif ($line !~ /:/ and defined $index->{$i-1}){
      # Here, we reconstitute multiline values (like status and action)
      chomp($line);
      $line =~ s/\s+/ /g;
      $errors->{$index->{$i-1}} .= $line;
    } elsif ($line =~ m/\s+[a-zA-Z0-9_\-]+\s+[A-Z]+\s+(?<read>\d+(\.\d+)?)(?<read_suffix>[KMT])?\s+(?<write>\d+(\.\d+)?)(?<write_suffix>[KMT])?\s+(?<cksum>\d+(\.\d+)?)(?<cksum_suffix>[KMT])?/){
      # And here, we count the number of read, write and checksum errors
      $errors->{read_errors}  += convert_suffix($+{'read'},$+{'read_suffix'});
      $errors->{write_errors} += convert_suffix($+{'write'},$+{'write_suffix'});
      $errors->{cksum_errors} += convert_suffix($+{'cksum'},$+{'cksum_suffix'});
    }
    $i++;
  }
  # Ensure evey item returns something
  $errors->{$_} ||= '' foreach (qw(scan action status errors));
  return $errors;
}

# Error counter can be suffixed. Apply this suffix to get raw error numbers
sub convert_suffix {
  my $val = shift;
  my $suf = shift;
  if (!$suf){
    return $val;
  } elsif ($suf eq 'K'){
    $val *= 1000;
  } elsif ($suf eq 'M') {
    $val *= 1000000;
  } elsif ($suf eq 'T') {
    $val *= 1000000000;
  }
  return $val;
}

sub get_zpool_stats {
  my $pool = shift;
  my $stats = {};
  return $stats unless (-e "/proc/spl/kstat/zfs/$pool/io");
  open STAT, "</proc/spl/kstat/zfs/$pool/io";
  while (<STAT>){
    if (m/^(?<nread>\d+)\s+(?<nwritten>\d+)\s+(?<reads>\d+)\s+(?<writes>\d+)\s+(?<wtime>\d+)\s+(?<wlentime>\d+)\s+(?<wupdate>\d+)\s+(?<rtime>\d+)\s+(?<rlentime>\d+)\s+(?<rupdate>\d+)\s+(?<wcnt>\d+)\s+(?<rcnt>\d+)/){
      $stats->{$_} = $+{$_} foreach (keys %+);
    }
  }
  close STAT;
  return $stats;
}