#!/usr/local/bin/expect -f
# 
# Written: Greg Baker (gregb@ifost.org.au)  23-Oct-2007
#
# This program logs in via ssh to the HP-UX systems given as command 
# line arguments, and resizes all filesystems in proportion to the
# standard deviation of their used capacity over time.
#
# ToDo:
#  - man page
#  - port to Linux
#  - support veritas volume manager
#  - port to Solaris
#  - smarter prompt comprehension
#  - cope with passwords and sites where you can't just log in as root
#  - report data to openview
#  - support clustered filesystems (symlink database files?)
#  - rewrite in Perl; use "observe" framework
  
set timeout -1

set arguments $argv

set collect yes
set fix yes
set force_fix no
set suggest yes
set verbose no
set show_activity yes

set dbdir /var/db
while {[string equal [string index [lindex $arguments 0] 0] -]} {
  set word [lindex $arguments 0]
  set arguments [lrange $arguments 1 end]
  case "$word" {
    --collect { set collect yes }
    --no-collect { set collect no }
    --fix { set fix yes }
    --no-fix { set fix no }
    --force-fix { set force_fix yes }
    --no-force-fix { set force_fix  no }
    --suggest { set suggest yes }
    --no-suggest { set suggest no }
    --verbose { set verbose yes }
    --show-activity { set show_activity yes }
    --hide-activity { set show_activity no }
    --db-dir { set dbdir [lindex $arguments 0] 
               set arguments [lrange $arguments 1 end] }
    * { puts stderr "Unknown argument: $word" ; exit }
  }
}

#####################################################################

proc sum_over_filesystem {which_array} {
 upvar $which_array a
 set sum 0
 foreach fs [array names a] { set sum [expr $sum + $a($fs)] }
 return $sum
}

set prompt "# "
# ToDo: be smarter in figuring out what the prompt is.

foreach hostname $arguments {
 if {!$show_activity} {log_user 0}
 spawn ssh -l root $hostname
 # expect password here perhaps?
 # cope with being an ordinary user and su later? Or use sudo as necessary?
 expect {
  "$prompt" {}
  "Are you sure you want to continue connecting (yes/no)? " { send "yes\n" ; exp_continue }
 }


 set volumegroups {}
 send "vgdisplay\n"

 expect {
  -re {VG Name */dev/([a-zA-Z0-9_-]*)} { 
     set vg $expect_out(1,string)
     lappend volumegroups $vg
     exp_continue
  }
  -re {PE Size \(Mbytes\) *([0-9][0-9]*)} { 
          set extent_size($vg) $expect_out(1,string)
          if {$verbose} { puts "Extent size is $extent_size($vg)" }
          set granularity($vg) [expr $extent_size($vg) * 1024]
          exp_continue
  }
  -re {Free PE *([0-9][0-9]*)} {
          set unallocated_space($vg) [expr $expect_out(1,string)  / 2]
          # the unallocated space is unmirrored... Naively we assume
          # everything is always mirrored.
          exp_continue
  }
  "$prompt" {

  }
 }

 if {$verbose} { puts "Unallocated spaces: [array get unallocated_space]" }
 if {[llength $volumegroups] == 0} {
  puts stderr "I tried to get the volume group information, but failed."
  exit 1
 }

 foreach vg $volumegroups {
  unset -nocomplain chunksavail totalavail stddevshares sharesize 
  unset -nocomplain lvextents_unused_by_filesystem currentavail stddevs
  unset -nocomplain can_shrink delta newlvolsize lvolsize new_fs_size
  unset -nocomplain currentsizes lvol 


   if {$verbose} { puts "bdf" }
   # ToDo: run uname and check what operating system it is. If it is HP-UX
   # then we can use bdf, otherwise use df. 

   send "bdf\n"
   expect {
    -re "/dev/$vg/(\[a-zA-Z0-9\]*) *(\[0-9\]*) *(\[0-9\]*) *(\[0-9\]*) *(\[0-9\]*)% (/\[a-zA-Z/0-9_-\]*)" {
     set filesystem $expect_out(6,string)
     set lvol($filesystem) $expect_out(1,string)
     set currentsizes($filesystem) $expect_out(2,string)
     set currentused($filesystem) $expect_out(3,string)
     set currentavail($filesystem) $expect_out(4,string)
     set lvextents_unused_by_filesystem($filesystem) [expr $currentavail($filesystem)/$granularity($vg)]
     set can_shrink($filesystem) yes
     exp_continue
    } 
    $prompt {
      if {$verbose} { puts "Got all the filesystems."}
    }
  }


  set frozen {}
  foreach {fs lv} [array get lvol] {
    if {$verbose} {puts "lvdisplay /dev/$vg/$lv"}
    send "lvdisplay /dev/$vg/$lv\n"
    expect "Logical volumes ---"

    expect {
      -re {Mirror copies *([012])} {
               set mirrors($fs) $expect_out(1,string)
             exp_continue
      }
      -re {LV Size \(Mbytes\) *([0-9]*)} { 
             set lvolsize($lv) $expect_out(1,string) 
             exp_continue 
      }
      -re {Current LE *([0-9]*)} { 
             set lvolextents($lv) $expect_out(1,string)
             exp_continue 
      }
      -re "Allocation *strict/contiguous" { 
             # one day I should cope with other continguous things
             lappend frozen $fs 
             exp_continue 
      }
      $prompt {  }
    }
  }

  foreach fs $frozen {
   set lv $lvol($fs)
   unset -nocomplain mirrors($fs)
   unset -nocomplain lvolsize($lv)
   unset -nocomplain lvolextents($lv) 
   unset -nocomplain lvol($fs)
   unset -nocomplain currentsizes($fs)
   unset -nocomplain currentused($fs)
   unset -nocomplain currentavail($fs)
   unset -nocomplain lvextents_unused_by_filesystem($fs)
   unset -nocomplain can_shrink($fs)
   if {$verbose} {puts "Can't work with $fs -- it wants contiguous allocation."}
  }

  # Make sure that there is some data.
  foreach fs [array names currentused] {
    set sample_count($fs) 0
    set sample_sum($fs) 0
    set sample_sum_of_squares($fs) 0
  }

  # Re-read old data if possible. Really should check for
  # sanity in hostnames and volume group names, otherwise
  # there's a small security problem here.
  if {[file exists "$dbdir/$hostname-$vg.txt"]} {
    set db [open "$dbdir/$hostname-$vg.txt" r]
    while {[gets $db lineitem] != -1} {
      set fs [lindex $lineitem 3]
      set sample_count($fs) [lindex $lineitem 0]
      set sample_sum($fs) [lindex $lineitem 1]
      set sample_sum_of_squares($fs) [lindex $lineitem 2]
      # we don't bother the stddev (last field in the line)
    }
    close $db
  }

  foreach fs [array names currentused] {
    incr sample_count($fs)
    set sample_sum($fs) [expr $sample_sum($fs) + 1.0 * $currentused($fs)] 
    set sample_sum_of_squares($fs) [expr (1.0 * $sample_sum_of_squares($fs)) + (1.0 * $currentused($fs) * 1.0 * $currentused($fs))] 
  }
  foreach fs [array names currentused] {
    set N $sample_count($fs)
    set SIGMA_X_SQ $sample_sum_of_squares($fs)
    set X_BAR [expr $sample_sum($fs) * 1.0 / $N]
    set tau [expr $SIGMA_X_SQ - $N * $X_BAR * $X_BAR]
    # watch out for loss of precision...
    if {$tau < 0.0} { set tau 0.0 }
    set stddevs($fs) [expr sqrt($tau/$N)]
  }
  # write it back again...
  if {!$collect} {
    file mkdir $dbdir
    set db [open "$dbdir/$hostname-$vg.temp" w]
    foreach fs [array names sample_count] {
      puts $db [list $sample_count($fs) $sample_sum($fs) $sample_sum_of_squares($fs) $fs $stddevs($fs)]
      if {$sample_count($fs) < 60} {
         if {$fix && !$force_fix} { 
            set fix no 
            if {$verbose} { puts "$fs only has $sample_count($fs) samples. This is not enough to calculate stddev. Pretending --no-fix was specified." }
         }
      }
    } 
    close $db
    file rename -force $dbdir/$hostname-$vg.temp $dbdir/$hostname-$vg.txt
  }

  foreach fs [array names stddevs] {
    # so that unchanging filesystems at least get some kind of share
    if {$stddevs($fs) < 0.001} { set $stddevs($fs) 1.0 }
  } 
 

  set all_shrinking_ok no
  while {!$all_shrinking_ok} {
    set all_shrinking_ok yes
    set chunksavail [expr $unallocated_space($vg)+[sum_over_filesystem lvextents_unused_by_filesystem]]
    set totalavail [sum_over_filesystem currentavail]
    set stddevshares [sum_over_filesystem stddevs]
    set sharesize [expr int($chunksavail / $stddevshares)]
    if {$verbose} {
      puts "Shrink capabilities: [array get can_shrink]"
      puts "It looks like we have $chunksavail logical extents to work with."
      puts "(they are currently spread out like this  [array get lvextents_unused_by_filesystem])"
      puts "Based around the following stddevs [array get stddevs]"
      puts "I should allocate $sharesize Mbytes per stddev."
      puts "Currently they are: [array get lvolsize]"
    }
 
    foreach filesystem [array names currentsizes] {
      set s $stddevs($filesystem)
      if {$can_shrink($filesystem)} {
        set desired_lvextents_unused [expr 1 + int(($s * $sharesize) / $extent_size($vg))] 
        set delta($filesystem) [expr $desired_lvextents_unused - $lvextents_unused_by_filesystem($filesystem)]
      } else {
        set desired_lvextents_unused 0
        set delta($filesystem) 0
      }
      set mylvol $lvol($filesystem)
      set newlvolsize($filesystem) [expr $lvolsize($mylvol) + $delta($filesystem)]
      set new_fs_size($filesystem) [expr $newlvolsize($filesystem)]
    }
    if {$verbose} {
      puts "I'm going to set these filesystem sizes (in Mb): [array get new_fs_size]"
    }

    if {$suggest} {
      puts "Filesystem suggested layout for $hostname, volume group $vg"
      foreach fs [array names delta] {
        puts "$fs ($lvol($fs)) should be: $new_fs_size($fs) Megabytes."
      }
      puts ""
      if {!$fix} { break }
    }

    foreach fs [array names delta] {
      if {$verbose} { puts "# Thinking about $fs, which has a delta of $delta($fs)" }
      if {$fix && ($delta($fs) < 0)} {
         send "fsadm -F vxfs -b $new_fs_size($fs)M $fs\n"
         expect {
           "size will be reduced" { exp_continue }
           "expand" { puts stderr "Faulty logic somewhere." ; exit 1 }
           "cannot shrink" { 
               set all_shrinking_ok no
               set lvextents_unused_by_filesystem($fs) [expr int($lvextents_unused_by_filesystem($fs)/2)]
               if {$lvextents_unused_by_filesystem($fs) < 1} {
                   set can_shrink($fs) no
               }
               if {$verbose} { puts " ... cannot shrink $fs as much as I would like. Have to recalculate based around that." }
	       expect "$prompt"
               break
           }
           "$prompt" { if {$verbose} { puts " ... $fs shrunk OK " }  }
         }
      }
    }
  }

  if {$fix} {
    foreach fs [array names delta] {
      if {$delta($fs) < 0} { 
        send "lvreduce -f -L $new_fs_size($fs) /dev/$vg/$lvol($fs)\n"
        expect $prompt
      }
    }

    foreach fs [array names delta] {
      if {$delta($fs) > 0} {
        send "lvextend -L $new_fs_size($fs) /dev/$vg/$lvol($fs)\n"
        expect $prompt
        send "fsadm -F vxfs -b $new_fs_size($fs)M $fs\n"
        expect $prompt
      }
    }
  } else {
    if {$verbose} { puts "Not attempting to fix."  }
  }
 }
 send "exit\n"
 expect eof
}