require 'sshkit/dsl'

class Router < CloudstackNagios::Base

  desc "memory HOST", "check memory on host"
  def memory
    begin
      host = systemvm_host
      free_output = ""
      on host do |h|
        free_output = capture(:free)
      end
      values = free_output.scan(/\d+/)
      total = values[0].to_i
      free = values[2].to_i
      free_b = values[7].to_i
      data = check_data(total, total - free_b, options[:warning], options[:critical])
      puts "MEMORY #{RETURN_CODES[data[0]]} - usage = #{data[1]}% (#{((total - free_b)/1024.0).round(0)}/#{(total/1024.0).round(0)}MB) | \
            usage=#{data[1]}% total=#{total}M free=#{free}M free_wo_buffers=#{free_b}M".gsub(/\s+/, " ")
      exit data[0]
    rescue => e
      exit_with_failure(e)
    end
  end

  desc "cpu", "check memory on host"
  def cpu
    begin
      host = systemvm_host
      mpstat_output = ""
      on host do |h|
        mpstat_output = capture(:mpstat)
      end
      values = mpstat_output.scan(/\d+\.\d+/)
      usage = 100 - values[-1].to_f
      data = check_data(100, usage, options[:warning], options[:critical])
      puts "CPU #{RETURN_CODES[data[0]]} - usage = #{data[1]}% | usage=#{data[1]}%"
      exit data[0]
    rescue => e
      exit_with_failure(e)
    end
  end

  desc "rootfs_rw", "check if the rootfs is read/writeable on host"
  def rootfs_rw
    begin
      host = systemvm_host
      test_file = '/rootdiskcheck.txt'
      rootfs_rw = false
      on host do |h|
        rootfs_rw = execute(:touch, test_file)
        execute(:rm, '-f', test_file)
      end
    rescue SSHKit::Command::Failed
      rootfs_rw = false
    rescue => e
      exit_with_failure(e)
    end
    status = rootfs_rw ? 0 : 2
    puts "ROOTFS_RW #{rootfs_rw ? 'OK - rootfs writeable' : 'CRITICAL - rootfs NOT writeable'}"
    exit status
  end

  desc "disk_usage", "check the disk space usage of the root volume"
  option :partition, desc: "The partition to check", default: '/', aliases: '-P'
  def disk_usage
    begin
      host = systemvm_host
      partition = options[:partition]
      proc_out = ""
      on host do |h|
        proc_out = capture(:df, '-l', partition)
      end
      match = proc_out.match(/.*\s(\d+)%\s.*/)
      if match
        usage = match[1]
        data = check_data(100, usage, options[:warning], options[:critical])
        puts "DISK_USAGE #{RETURN_CODES[data[0]]} (Partition #{options[:partition]}) - usage = #{data[1]}% | usage=#{data[1]}%"
        exit data[0]
      else
        puts "DISK_USAGE UNKNOWN"
      end
    rescue => e
      exit_with_failure(e)
    end
  end

  desc "network", "check network usage on host"
  option :interface,
      desc: 'network interface to probe',
      default: 'eth0',
      aliases: '-i'
  option :if_speed,
      desc: 'network interface speed in bits per second',
      type: :numeric,
      default: 1000000,
      aliases: '-s'
  def network
    begin
      host = systemvm_host
      stats_path = "/sys/class/net/#{options[:interface]}/statistics"
      rx_bytes, tx_bytes = ""
      on host do |h|
        rx_bytes = capture("cat #{stats_path}/rx_bytes;sleep 1;cat #{stats_path}/rx_bytes").lines.to_a
        tx_bytes = capture("cat #{stats_path}/tx_bytes;sleep 1;cat #{stats_path}/tx_bytes").lines.to_a
      end
      rbps = (rx_bytes[1].to_i - rx_bytes[0].to_i) * 8
      tbps = (tx_bytes[1].to_i - tx_bytes[0].to_i) * 8
      data = check_data(options[:if_speed], rbps, options[:warning], options[:critical])
      puts "NETWORK #{RETURN_CODES[data[0]]} - usage = #{data[1]}% | usage=#{data[1]}% rxbps=#{rbps.round(0)} txbps=#{tbps.round(0)}"
      exit data[0]
    rescue => e
      exit_with_failure(e)
    end
  end

  desc "conntrack_connections", "check the number of conntrack connections"
  def conntrack_connections
    begin
      host = systemvm_host
      default_max = 1000000
      netfilter_path = "/proc/sys/net/netfilter/"
      current, max = 0
      on host do |h|
        max     = capture("cat #{netfilter_path}nf_conntrack_max").to_i
        current = capture("cat #{netfilter_path}nf_conntrack_count").to_i
      end
      if max < default_max
        on host do |h|
          execute :echo, "#{default_max} > #{netfilter_path}nf_conntrack_max"
        end
      end
      data = check_data(max, current, options[:warning], options[:critical])
      puts "CONNTRACK_CONNECTIONS #{RETURN_CODES[data[0]]} - usage = #{data[1]}% (#{current.round(0)}/#{max.round(0)}) | usage=#{data[1]}% current=#{current.round(0)} max=#{max.round(0)}"
      exit data[0]
    rescue => e
      exit_with_failure(e)
    end
  end

  desc "active_ftp", "make sure conntrack_ftp and nf_nat_ftp modules are loaded"
  def active_ftp
    begin
      host = systemvm_host
      active_ftp_enabled = false
      modules = %w(nf_conntrack_ftp, nf_nat_ftp)
      on host do |h|
        lsmod = capture(:lsmod)
        active_ftp_enabled = lsmod.include?('nf_conntrack_ftp') &&
          lsmod.include?('nf_nat_ftp')
        unless active_ftp_enabled
          # load the modules in the kernel
          execute(:modprobe, 'nf_conntrack_ftp')
          execute(:modprobe, 'nf_nat_ftp')
          # load the modules at next server boot
          execute(:echo, '"nf_conntrack_ftp" >> /etc/modules')
          execute(:echo, '"nf_nat_ftp" >> /etc/modules')
          active_ftp_enabled = true
        end
      end
    rescue SSHKit::Command::Failed
      active_ftp_enabled = false
    rescue => e
      exit_with_failure(e)
    end
    status = active_ftp_enabled ? 0 : 2
    puts "ACTIVE_FTP #{active_ftp_enabled ? 'OK - active_ftp enabled' : 'CRITICAL - active_ftp NOT enabled'}"
    exit status
  end

  no_commands do

    def systemvm_host
      unless options[:host]
        say "Error: --host/-H option is required for this check.", :red
        exit 1
      end
      host = SSHKit::Host.new("root@#{options[:host]}")
      host.ssh_options = sshoptions(options[:ssh_key])
      host.port = options[:ssh_port]
      host
    end

  end

end