#!/usr/bin/env perl
use warnings;
use strict;  

use version ; our $VERSION = 'v2.1.0';

# Use smartctl, lsblk and geom to get disk information
# Author: R. W. Rodolico
# Date:   2018-04-02
#
# will decode lsblk first. Then use geom to get information on BSD systems
# finally, smartctl will be used to gather additional information.
#
# Revision History
#
# 20200224 RWR v2.1
# large re-organization of code and inclusion to get basic information from lsblk if it is on the system
#

# find our location and use it for searching for libraries. library.pm must be in the same directory as the calling script
# or, if run interactively, in the parent of the modules
BEGIN {
   use FindBin;
   use File::Spec;
   # prepend the bin directory and its parent
   use lib File::Spec->catdir($FindBin::Bin), File::Spec->catdir("$FindBin::Bin/..");
   eval( 'use library;' );
   die sprintf( "Could not find library.pm in %s, INC is %s\n", __FILE__, join( "\n", @INC ) ) if $@;
}

#####
##### Change these to match your needs
#####

# Make this a list of all the modules we are going to use. You can replace undef with the version you need, if you like
my $modulesList = {
        'Data::Dumper'     => undef,
   };

# hash of commands that are needed for the system. key is the name of the command and, in some cases, the value will become
# the full path (from which or where)

# WARNING: This will only work on FreeBSD until the system is updated. We only work with geom
my $commandsList = {
        'smartctl' => undef,
#        'geom' => undef,
#        'lsblk' => undef,
   };

# list of operating systems this module can be used on.
my $osList = {
#         'mswin32' => undef,
         'freebsd' => undef,
         'linux'   => undef,
   };

# the category the return data should go into. See sysinfo for a list
my $CATEGORY = 'attributes';

#####
##### End of required
#####

# some variables needed for our system
my $errorPrepend = 'error: in ' . __FILE__; # this is prepended to any error messages
my @out; # temporary location for each line of output

# Try to load the modules we need. If we can not, then make a list of missing modules for error message.
for my $module ( keys %$modulesList ) {
   eval ( "use $module;" );
   push @out, "$errorPrepend Could not load $module" if $@;
}

if ( ! @out && ! checkOS ( $osList ) ) { # check if we are on an acceptible operating system
    push @out, "$errorPrepend Invalid Operating System";
}
if ( !@out && ! validCommandOnSystem ( $commandsList ) ) {
   push @out, "$errorPrepend Can not find some commands needed";
}
if ( !@out ) { # we made it, we have everything, so do the processing
   #####
   ##### Your code starts here. Remember to push all output onto @out
   #####
   
   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;
      }
      
   }

   # 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();

   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}} ) {
         push @out,  "$CATEGORY\t$drive $key\t$driveDefinitions{$drive}{$key}" if defined $driveDefinitions{$drive}{$key};
      }
   }

   
   #####
   ##### Your code ends here.
   #####
}

# If we are testing from the command line (caller is undef), print the results for debugging
print join( "\n", @out ) . "\n" unless caller;
# called by do, which has a value of the last assignment made, so make the assignment. The equivilent of a return
my $return = join( "\n", @out );