#!/usr/bin/env perl
use warnings;
use strict;
use Data::Dumper;

# Description: Use smartctl, lsblk and geom to get disk information
# will decode lsblk first. Then use geom to get information on BSD systems
# finally, smartctl will be used to gather additional information.

# 20200224 RWR v2.1
# large re-organization of code and inclusion to get basic information from lsblk if it is on the system

our $VERSION = '2.1';

# R. W. Rodolico
# grabs smart readings from all drives in system
# note that in the case of some hardware RAID controller, will list the components of a "drive" and the drive itself,
# so you may have up to twice the number of entries than you have physical drives, for example, if you have a megaraid
# controller where each physical drive is also a single logical drive.

# find our location and use it for searching for libraries
BEGIN {
   use FindBin;
   use File::Spec;
   use lib File::Spec->catdir($FindBin::Bin);
   use library;
}

exit 0 unless checkDate( 'm' ); # run this only on first of month

my %driveDefinitions; # this will be a global that everything will put info into

my %ignoreDriveTypes = ( # a list of drive "types" used by virtuals that we just ignore.
         'VBOX HARDDISK' => 1,
         );


# routine used by geom to process one block at a time.
# first parameter is a pointer to a hash to populate, the rest are
# considered to be a splice of an array that contains only one
# block. Returns the name of the device
sub doGeomBlock {
   my $hashPointer = shift;
#   die join( "\n", @_ ) . "\n";
   while ( my $line = shift ) {
      if ( $line ) {
         my ($key, $value) = split( ':', $line );
         if ( $key =~ m/^\d+\.\s+(.*)$/ ) {
            $key = $1;
         }
         $key = &trim( $key );
         $value = &trim( $value );
         $hashPointer->{$key} = $value;
      }
   }
#   die Dumper( $hashPointer );
   return $hashPointer->{'Name'} ? $hashPointer->{'Name'} : 'Unknown';
}

# grab data from geom command (BSD) and import parts of it into driveDefinitions
sub geom {
   my $line = 0;
   my $startBlock = 0;
   my $endBlock = 0;
   my @report = `geom disk list`;
   chomp @report;

   while ( $line < scalar( @report ) ) {
      #print "Working on $line\n";
      while ( $line < scalar( @report ) && $report[$line] !~ m/^Geom name:\s+(.*)$/ ) {
         $line++;
      }
      $endBlock = $line - 1;
      if ( $endBlock > $startBlock ) {
         my $thisDrive = {};
         my $key = &doGeomBlock( $thisDrive, @report[$startBlock..$endBlock] );
         # die "$key\n" . Dumper( $thisDrive );
         $key = '/dev/' . $key;
         $driveDefinitions{$key}{'Capacity'} = $thisDrive->{'Mediasize'} if defined $thisDrive->{'Mediasize'};
         if ( defined( $driveDefinitions{$key}{'Capacity'} ) && $driveDefinitions{$key}{'Capacity'} =~ m/^(\d+)/ ) {
            $driveDefinitions{$key}{'Capacity'} = $1;
         }
         $driveDefinitions{$key}{'Model'} = $thisDrive->{'descr'} if defined $thisDrive->{'descr'};
         $driveDefinitions{$key}{'Serial'} = $thisDrive->{'ident'} if defined $thisDrive->{'ident'};
         $driveDefinitions{$key}{'Sector Size'} = $thisDrive->{'Sectorsize'} if defined $thisDrive->{'Sectorsize'};
         $driveDefinitions{$key}{'Rotation'} = $thisDrive->{'rotationrate'} if defined $thisDrive->{'rotationrate'};

         $startBlock = $line;
#         die Dumper( $thisDrive );
      }
      $line++;
   }

}



# acquires information using lsblk, if it is available
# uses global %driveDefinitions to store the results

sub lsblk {
   eval ( 'use JSON qw( decode_json );' );
   if ( $@ ) {
      warn "Could not load JSON library\n";    
      return;
   }
   
   my $output = qx'lsblk -bdJO 2>/dev/null';
   # older versions do not have the O option, so we'll run it without
   $output = qx'lsblk -bdJ 2>/dev/null' if $?;
   my $drives = decode_json( join( '', $output ) );
   $drives = $drives->{'blockdevices'};
   while ( my $thisDrive = shift @{$drives} ) {
      if ( $thisDrive->{'type'} eq 'disk' ) {
         my $key = '/dev/' . $thisDrive->{'name'};
         $driveDefinitions{$key}{'Capacity'} = $thisDrive->{'size'};
         $driveDefinitions{$key}{'Model'} = $thisDrive->{'model'} if defined $thisDrive->{'model'};
         $driveDefinitions{$key}{'Serial'} = $thisDrive->{'serial'} if defined $thisDrive->{'serial'};
      }
   }
}


sub getSmartInformationReport {
   my ($drive, $type) = @_;
   $type = '' if ( $type =~ m/scsi/ );
   my @report = `smartctl -i $drive $type`;
   chomp @report;
   my %reportHash;
   for ( my $i = 0; $i < @report; $i++ ) {
      if ( $report[$i] =~ m/^(.*):(.*)$/ ) {
         $reportHash{$1} = trim($2);
      }
   }
   return \%reportHash;
} # getSmartInformationReport

sub getSmartAttributeReport {
   my ($drive, $type) = @_;
   $type = '' if ( ! defined( $type ) || $type =~ m/scsi/ );
   my @report = `smartctl -A $drive $type`;
   chomp @report;
   my %reportHash;
   my %headers;
   # bypass all the header information
   my $i;
   for ( $i = 0; $i < @report && $report[$i] !~ m/^ID#/; $i++ ) {}
   if ( $i < @report ) { # did we get an actual report? some drives will not give us one
      my $char = 0;
      while ( $char < length($report[$i]) ) {
         substr( $report[$i],$char ) =~ m/^([^ ]+\s*)/;
         my $header = $1;
         my $start = $char;
         my $length = length($header);
         if ( $header = &trim( $header ) ) {
            $headers{$header}{'start'} = $start;
            $headers{$header}{'length'} = $length-1;
         }
         $char += $length;
      }
      while ( ++$i < @report ) {
         last unless $report[$i];
         my $id = &trim(substr( $report[$i], $headers{'ID#'}{'start'}, $headers{'ID#'}{'length'} ));
         my $name = &trim(substr( $report[$i], $headers{'ATTRIBUTE_NAME'}{'start'}, $headers{'ATTRIBUTE_NAME'}{'length'} ));
         my $value = &trim(substr( $report[$i], $headers{'RAW_VALUE'}{'start'} ));
         $reportHash{$id}{'value'} = $value;
         $reportHash{$id}{'name'} = $name;
      }
   }
   #print Dumper( \%reportHash ); die;
   return \%reportHash;
}


sub getAttributes {
   my ($drive,$type,$sectorSize) = @_;

   my $report = &getSmartAttributeReport( $drive, $type );

   # first let's get total disk writes
   if ( defined( $report->{'241'} ) && defined( $sectorSize ) ) {
      $sectorSize =~ m/^(\d+)/;
      $driveDefinitions{$drive}{'writes'} = $report->{'241'} * $sectorSize;
   }
   # find total uptime
   if ( defined( $report->{'9'} ) ) {
      $report->{'9'}->{'value'} =~ m/^(\d+)/;
      $driveDefinitions{$drive}{'uptime'} = $1;
   }
}

sub getInformation {
   my ($drive, $type) = @_;

   my $report = &getSmartInformationReport( $drive, $type );

   my %info;
   my %keys = ( 
                  'Model Family'  => { 
                                          'tag' => 'Make',
                                          'regex' => '(.*)'
                                      },
                  'Device Model'  => { 
                                          'tag' => 'Model',
                                          'regex' => '(.*)'
                                      },
                  'Serial number' => { 
                                          'tag' => 'Serial',
                                          'regex' => '(.*)'
                                      },
                  'Serial Number' => { 
                                          'tag' => 'Serial',
                                          'regex' => '(.*)'
                                      },
                  'User Capacity' => { 
                                          'tag' => 'Capacity',
                                          'regex' => '([0-9,]+)'
                                      },
                  'Logical block size' =>{ 
                                          'tag' => 'Sector Size',
                                          'regex' => '(\d+)'
                                      },
                  'Sector Size'   => { 
                                          'tag' => 'Sector Size',
                                          'regex' => '(\d+)'
                                      },
                  'Rotation Rate' => { 
                                          'tag' => 'Rotation',
                                          'regex' => '(.*)'
                                      },
               );
   foreach my $key ( keys %keys ) {
      if ( defined( $report->{$key} ) && $report->{$key} =~ m/$keys{$key}->{'regex'}/ ) {
         $driveDefinitions{$drive}{$keys{$key}->{'tag'}} = $1;
      }
   }
}

sub smartctl {
   # Get all the drives on the system
   my %allDrives;
   eval{ 
      %allDrives = map { $_ =~ '(^[a-z0-9/]+)\s+(.*)\#'; ($1,$2) } `smartctl --scan`; 
   };
   return if ( $@ ); # we failed to get anything back
      

   # output the drives and information as tab delimited,
   foreach my $thisDrive ( sort keys %allDrives ) {
      $driveDefinitions{$thisDrive}{'type'} = $allDrives{$thisDrive};
      &getInformation( $thisDrive, $allDrives{$thisDrive} );
      #print Dumper( $info ); die;
      my $attributes = &getAttributes( $thisDrive, $allDrives{$thisDrive}, $driveDefinitions{$thisDrive}{'Sector Size'} );
      #print Dumper( $attributes ); die;
      #print Dumper( $info ); die;
   }
   
}



# category we will use for all values found
# see sysinfo for a list of valid categories
my $CATEGORY = 'attributes';

# run the commands necessary to do whatever you want to do
# The library entered above has some helper routines
# validCommandOnSystem -- passed a name, returns the fully qualified path or '' if it does not exist
# cleanUp - passed a delimiter and a string, does the following (delimiter can be '')
#           chomps the string (removes trailing newlines)
#           removes all text BEFORE the delimiter, the delimiter, and any whitespace
#           thus, the string 'xxI Am x  a weird string' with a newline will become
#           'a weird string' with no newline

# now, return the tab delimited output (to STDOUT). $CATEGORY is the first
# item, name as recognized by sysinfo is the second and the value is
# the last one. For multiple entries, place on separate lines (ie, newline separated)

# check for commands we want to run
my %commands = ( 
                  'smartctl' => '',
                  'lsblk' => '',
                  'geom' => '',
               );

# check the commands for validity
foreach my $command ( keys %commands ) {
   $commands{$command} = &validCommandOnSystem( $command );
}

# bail if we don't have the commands we need
exit 1 unless $commands{'smartctl'} || $commands{'lsblk'} || $commands{'geom'};

# first, get basic information using lsblk for linux systems
&lsblk() if $commands{'lsblk'};
# now, try geom for bsd systems
&geom() if $commands{'geom'};

#die Dumper( \%driveDefinitions );
# finally, populate whatever using smartctl if it is on system.
&smartctl() if $commands{'smartctl'};

for my $drive ( sort keys %driveDefinitions ) {
   # don't print iSCSI definitions
   next if defined( $driveDefinitions{$drive}{'Transport protocol'} ) && $driveDefinitions{$drive}{'Transport protocol'} eq 'ISCSI';
   #also, blow off our ignored types
   next if ( defined( $driveDefinitions{$drive}{'Model'} ) && defined( $ignoreDriveTypes{ $driveDefinitions{$drive}{'Model'} }  ) );
   
   # remove comma's from capacity
   $driveDefinitions{$drive}{'Capacity'} =~ s/,//g if $driveDefinitions{$drive}{'Capacity'};
   foreach my $key ( sort keys %{$driveDefinitions{$drive}} ) {
      print "$CATEGORY\t$drive $key\t$driveDefinitions{$drive}{$key}\n" if defined $driveDefinitions{$drive}{$key};
   }
}


exit 0;