#! /usr/bin/env perl

# getDriveInfo
# 
# Script to get information on drives for all Unix systems. It will retrieve 
# manufacturing information, SMART information and, with software RAID or 
# ZFS file systems, membership.
#
# Also looks for a file, /etc/drive_bays to determine which physical bay a
# drive is installed in (tab separated, first line header's 'bay' and 'serial'
# followed by the data).
#
# generates a tab separated table of all values retrieved
#
# Copyright (c) 2024, Daily Data, Inc.
# 
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 
#   Redistributions of source code must retain the above copyright notice,
#   this list of conditions and the following disclaimer.
#   Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
# 
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR 
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
#
# v0.1 RWR 20240708
# Added ability to read tab separated list of bays and serial numbers 
# in /etc/drive_bays to populate the physical bay's a drive is in. Format
# (\t means "tab"):
# bay\tserial
# 0\tserial number of drive 0
# 1\tserial number of drive in bay 1
# bay can be any string, so in one case we use 'internal' as the bay number

use strict;
use warnings;

use Data::Dumper;

my $TESTING = 0;

my $driveDefinitions; # hashref is global that everything will put info into

my $ignoreDriveTypes = { # a list of fake drive "types" that we just ignore.
         'VBOX HARDDISK' => 1, # virtualbox
         'CTLDISK' => 1,   # iSCSI drive
         };

# check for commands we want to run, so some of the must be on the
# system
my $commands = {
                  'smartctl' => '',
                  'lsblk' => '',
                  'geom' => '',
                  'zpool' => '',
                  'mdadm' => '',
                  'glabel' => '',
                  
               };
               

# remove leading and trailing spaces from string
sub trim {
 my $value = shift;
 $value =~ s/^\s+|\s+$//g if $value;
 return $value;
}

# 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
# will also look for single and double quotes surrounding entire string and remove them
# if they exist

sub cleanUp {
   my ($delimiter, $text) = @_;
   chomp $text;
   if ( $delimiter && $text =~ m/[^$delimiter]*$delimiter\s*(.*)/ ) {
      $text = $1;
   }
   if ( $text =~ m/^'(.*)'$/ ) {
      $text = $1;
   }
   if ( $text =~ m/^"(.*)"$/ ) {
      $text = $1;
   }
   return $text;
}


# checks if a command is valid on this system. If so, returns the full path to it
# else, returns empty string.
sub validCommandOnSystem {
   my $command = shift;
   $command = `which $command 2> /dev/null`;
   chomp $command;
   return -x $command ? $command : '';
}

# returns the operation system, all lower case
sub getOperatingSystem {
   return lc &cleanUp('', qx(uname -s));
}   


# acquires information using lsblk (Linux), 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;
   if ( $TESTING ) {
      $output = `cat lsblk`; 
   } else {
      $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'} = defined $thisDrive->{'model'} ? $thisDrive->{'model'} : '';
         $driveDefinitions->{$key}->{'serial'} = defined( $thisDrive->{'serial'} ) ? $thisDrive->{'serial'} : '' ;
         $driveDefinitions->{$key}->{'sector'} = '';
      }
   }
}

# Acquires disk information if it available (FreeBSD)
# uses global %driveDefinitions to store the results
sub geom {
   my %disks;
   my @data = `geom disk list`;
   
   my %conversion = (
      'ident' => 'serial',
      'descr' => 'model',
      'Mediasize' => 'capacity',
      'Sectorsize' => 'sector'
      );
   
   my $currentDisk = '';
   for ( my $i = 0; $i < @data; $i++ ) {
      if ( $data[$i] =~ m/^Geom name:\s+(.*)/ ) {
         $currentDisk = &fixDisk($1);
      } elsif ( $data[$i] =~ m/^\s+([^:]+):\s+(.*)/ ) {
        $driveDefinitions->{$currentDisk}->{$conversion{$1}} = $2 if $conversion{$1};
      }
   }
   return 1;
}

# determines what to call to get drive definitions
sub getDriveDefinitions {
   if ( $commands->{'lsblk'} ) { # must be a Linux system
      &lsblk();
   } elsif ( $commands->{'geom'} ) { # FreeBSD
      &geom();
   } else {
      die "Could not locate lsblk or geom to get list of disks, aborting\n";
   }
   return 1;
}

# Actual call to smartctl
sub getSmartInformationReport {
   my ($drive, $type) = @_;
   $type = '' if ( $type =~ m/scsi/ );
   my @report = `smartctl -j -a $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 getSmartInformation {
   my ( $drive,$type ) = @_;
   # my $report = &getSmartInformationReport( $drive, $type );
   $type = $type eq 'scsi' ? '' : "-t $type";
   my $report = &getJson( "smartctl -j $type -a $drive" );
   #print Dumper( $report ); die;
   # this is what we look for in the smart report
   $driveDefinitions->{$drive}->{'make'} = defined( $report->{'model_name'} ) ? $report->{'model_name'} : '';
   $driveDefinitions->{$drive}->{'model'} = defined( $report->{'model_family'} ) ? $report->{'model_family'} : '';
   $driveDefinitions->{$drive}->{'serial'} = defined( $report->{'serial_number'} ) ? $report->{'serial_number'} : '';
   $driveDefinitions->{$drive}->{'capacity'} = defined(    $report->{'user_capacity'}->{'bytes'} ) ?    $report->{'user_capacity'}->{'bytes'} : '';
   $driveDefinitions->{$drive}->{'sector'} = defined( $report->{'logical_block_size'} ) ? $report->{'logical_block_size'} : '';
   $driveDefinitions->{$drive}->{'rotation'} = defined( $report->{'rotation_rate'} ) ? $report->{'rotation_rate'} : '';

   $driveDefinitions->{$drive}->{'reboots'} = defined( $report->{'power_cycle_count'} ) ? $report->{'power_cycle_count'} : '';
   $driveDefinitions->{$drive}->{'temperature'} = defined( $report->{'temperature'}->{'current'} ) ? $report->{'temperature'}->{'current'} : '';
   $driveDefinitions->{$drive}->{'form_factor'} = defined( $report->{'form_factor'}->{'name'} ) ? $report->{'form_factor'}->{'name'} : '';
   $driveDefinitions->{$drive}->{'power_on_hours'} = defined( $report->{'power_on_time'}->{'hours'} ) ? $report->{'power_on_time'}->{'hours'} : '';

   my $toCheck = {
      '241' => 'lba_write',
      '233' => 'wearout',
      '9' => 'hours_on',
      '12' => 'power_cycle'
   };

   # make sure there are (blank) entries for everything
   foreach my $this ( keys %$toCheck ) {
      $driveDefinitions->{$drive}->{$toCheck->{$this}} = '';
   }
   # Now, fill in when we can find it
   foreach my $entry ( @{$report->{'ata_smart_attributes'}->{'table'}} ) {
      if ( defined( $toCheck->{$entry->{'id'}} ) ) {
         $driveDefinitions->{$drive}->{$toCheck->{$entry->{'id'}} . '_raw' } = $entry->{'raw'}->{'value'};
         $driveDefinitions->{$drive}->{$toCheck->{$entry->{'id'}} } = $entry->{'value'};
      }
   }


   return 1;
   my %keys = ( 
                  '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;
      }
   }
   return 1;
}

# reads a JSON command and decodes it
sub getJson {
   my $command = shift;
   eval ( 'use JSON qw( decode_json );' );
   if ( $@ ) {
      warn "Could not load JSON library when processing [$command]\n";    
      return;
   }
   my $output = qx/$command/;
   return decode_json( join( '', $output ) );
}


sub getSmart {
   # by running scan, we get the type of drive which we are looking at
   #my %allDrives = map { $_ =~ '(^[a-z0-9/]+)\s+(.*)\#'; ($1,$2) } `smartctl --scan`;
   my $drives = getJson( 'smartctl -j --scan' );
   $drives = $drives->{'devices'};
   my %allDrives;
   foreach my $thisDrive ( @$drives ) {
      $allDrives{$thisDrive->{'name'} } =  $thisDrive->{'type'} if $thisDrive->{'type'} ne 'atacam';
   }

   # Add smart information to the definitions,
   foreach my $thisDrive ( sort keys %allDrives ) {
      $driveDefinitions->{$thisDrive}->{'type'} = $allDrives{$thisDrive};
      &getSmartInformation( $thisDrive, $driveDefinitions->{$thisDrive}->{'type'} );
      $driveDefinitions->{$thisDrive}->{'capacity'} =~ s/,//g;
   }
   return 1;
}

sub mdadm {
   my @mdstat;
   if ( $TESTING ) {
      @mdstat = `cat mdstat`;
   } else {
      @mdstat = `cat /proc/mdstat`;
   }
   for ( my $line = 0; $line < @mdstat; $line++ ) {
      if ( $mdstat[$line] =~ m/^(md\d+)[\s:]+active (raid\d+)\s(.*)$/ ) { # this is the definition of an MD
         my $md = $1;
         my $raidLevel = $2;
         my @drives = split( /\s/, $3 );
         foreach my $currentDisk ( @drives ) {
            $currentDisk =~ m/^([a-z]+)(\d?)\[(\d+)\]/;
            my $disk = &fixDisk( $1 );
            my $component = $3;
            $driveDefinitions->{$disk}->{'vdev_type'} = $raidLevel;
            $driveDefinitions->{$disk}->{'pool'} = $md;
            $driveDefinitions->{$disk}->{'component'} = $3;
         }
      } # if it is a RAID member
   } # for $line
   return 1;
}

# some BSD systems will have partitions labeled before they
# are put into a zfs file system. This will build a translation
# for them
# returns 'label' => 'device'
# 'gptid/c73b7eae-f242-11ed-b0df-ecf4bbdad680' => 'da0p1'
sub glabel {
   my $labels;
   my @drives = `glabel status -s`;
   for ( my $i = 0; $i < @drives; $i++ ) {
      my @fields = split( /\s+/, $drives[$i] );
      $labels->{$fields[0]} = '/dev/' . $fields[2];
   }
   return $labels;
}

# ensure /dev/ in front of disk name   
sub fixDisk {
   my $disk = shift;
   my $partition = '';
   if ( $disk =~ m|^(/dev/)?([a-z]{2}\d+)p(\d+)$| ) { # BSD style
      $disk = "/dev/$2";
      $partition = $3;
   } elsif ( $disk =~ m|^(/dev/)?([a-z]{3})(\d+)$| ) { # linux style
      $disk = "/dev/$2";
      $partition = $3;
   }
   $disk = '/dev/' . $disk unless $disk =~ m/\/dev\//;
   return $disk
}

sub zfs {
   # grab a list of all pools and associated drives
   my @zpool = `zpool status`;
   # grab label to drive information
   my $labels = &glabel();
   my $i=0;
   while ( $i < @zpool ) {
      # bypass everything up to the pool: flag
      while ( $i < @zpool && $zpool[$i] !~ m/^\s*pool:\s*(\S+)/ ) {
         $i++;
      }
      last unless $i < @zpool;
      $zpool[$i] =~ m/^\s*pool:\s*(\S+)/;
      my $currentPool = $1;
      # bypass everything up to the headers
      while ( $i < @zpool && $zpool[$i] !~ m/^(\s+)NAME\s+STATE/ ) {
         $i++;
      }
      $zpool[$i] =~ m/^(\s+)NAME\s+STATE/;
      # the number of spaces preceding NAME is the number of spaces preceding the pool name
      my $poolNameIndent = $1;
      $i++;
      # The final loop will go through all the items which are in the pool configurations
      # and parse them out.
      my $currentComponent = '';
      my $poolType = '';
      my $poolTypeIndent = '';
      my $poolDeviceIndent = '';
      while ( $zpool[$i] =~ m/^$poolNameIndent(.*)/ ) {
         my $line = $1;
         if ( $line =~ m/^([a-zA-Z0-9_-]+)/ ) { # this is a pool name
            $currentComponent = $1;
         } elsif ( ! $poolTypeIndent && $line =~ m/^(\s+)([a-zA-Z0-9_-]+)/ ) {
            $poolTypeIndent = $1;
            $poolType = $2;
         } elsif ( $line =~ m/^$poolTypeIndent([a-zA-Z0-9_-]+)/ ) {
            $poolType = $1;
         } elsif ( ! $poolDeviceIndent && $line =~ m/^(\s+)([a-zA-Z0-9_\/-]+)/ ) {
            my $disk = &fixDisk( $labels->{$2} ? $labels->{$2} : $2 );
            $poolDeviceIndent = $1;
            $driveDefinitions->{$disk}->{'vdev_type'} = $poolType;
            $driveDefinitions->{$disk}->{'pool'} = $currentPool;
            $driveDefinitions->{$disk}->{'component'} = $currentComponent;
         } elsif ( $line =~ m/^$poolDeviceIndent([a-zA-Z0-9_\/-]+)/ ) {
            my $disk = &fixDisk($labels->{$1} ? $labels->{$1} : $1);
            $driveDefinitions->{$disk}->{'vdev_type'} = $poolType;
            $driveDefinitions->{$disk}->{'pool'} = $currentPool;
            $driveDefinitions->{$disk}->{'component'} = $currentComponent;
         }
         $i++;
      } # while
   } # while
   return 1;
}

sub getUsage {
   foreach my $drive ( keys %$driveDefinitions ) {
      $driveDefinitions->{$drive}->{'vdev_type'} = '';
      $driveDefinitions->{$drive}->{'pool'} = '';
      $driveDefinitions->{$drive}->{'component'} = '';
   }
   if ( $TESTING || $commands->{'mdadm'} ) {
      mdadm();
   } elsif ( $commands->{'zpool'} ) {
      zfs();
   }
      
   return 1;
}

# if /etc/drive_bays exists and is in proper format
# include the location of each drive
sub getDriveBays {
   my $driveBays = ();
   if ( -f '/etc/drive_bays' ) {
      open DATA, "</etc/drive_bays";
      my $line = <DATA>;
      while ( my $line = <DATA> ) {
         chomp $line;
         next unless $line;
         my ($bay, $serial) = split( "\t", $line );
         
         $driveBays->{$serial} = $bay;
      }
   }
   foreach my $drive ( keys %{$driveDefinitions} ) {
      $driveDefinitions->{$drive}->{'bay'} = defined( $driveBays->{$driveDefinitions->{$drive}->{'serial'}} ) ? $driveBays->{$driveDefinitions->{$drive}->{'serial'}} : '';
   }
   return 1;
}


sub report {
   my @headers;
   my @report;

   for my $drive ( sort keys %$driveDefinitions ) {
      my @line;
      # 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'};
      push @line, $drive;
      unless ( @headers ) {
         # for some reason, doing this the fast way (with keys) did not work. some of the keys were joined together
         # so, doing it this way.
         push @headers, 'drive';
         foreach my $header ( sort keys %{$driveDefinitions->{$drive}} ) {
            $header =~ s/[^a-z0-9_]//gi;
            push @headers, $header;
         }
         #@headers = sort keys %{$driveDefinitions->{$drive}} unless @headers;
         #unshift @headers,'drive';
      }
      foreach my $key ( sort keys %{$driveDefinitions->{$drive}} ) {
         push @line,$driveDefinitions->{$drive}->{$key};
      }
      push @report, join( "\t", @line );;
   }
   unshift @report,join ("\t",@headers);
   return join( "\n", @report ) . "\n";
   #return Dumper( $driveDefinitions );;
}

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

my $os = &getOperatingSystem();

# get the drives in the machine with minimal information
&getDriveDefinitions( $driveDefinitions );

# figure out what the drives are used for, ie mdadm or zfs or something
&getUsage( );

&getSmart();

&getDriveBays();

print &report( $driveDefinitions );

1;