#!/usr/bin/env python
##############################################################################
# Copyright (c) Members of the EGEE Collaboration. 2011.
# See http://www.eu-egee.org/partners/ for details on the copyright
# holders.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS
# OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##############################################################################
#
# NAME :        check_dpm_pool
#
# DESCRIPTION : Checks the pool freespace and status
#
# AUTHORS :     Alejandro.Alvarez.Ayllon@cern.ch
#
##############################################################################

import commands
import dpm
import os
import time
from lcgdmcommon import *

class check_dpm_pool:
  "Checks the DPM pools usage"
  __version__     = "0.0.1"
  __nagios_id__   = "DM-POOL"

  # Defaults
  DEFAULT_WARNING  = "5G"
  DEFAULT_CRITICAL = "1G"
  DEFAULT_UNITS     = None

  # Specific parameters, where key = short, value = long (i.e. {"h":"help", "C:":"command="})
  # getopt format. The long version will be the one passed even when the short is specified
  __additional_opts__ = {"w:": "warning=",
                         "c:": "critical=",
  			 "u:": "units=",
                         "p:": "pools=",
                         "O:": "VO="}

  # Specific usage information
  __usage__ = """
\t-w, --warning\tSets the warning limit for free space. It can be two values: pool. It accepts suffixes. (e.g. -w 100G). Default %s.
\t-c, --critical\tSets the critical limit for free space. It can be two values: pool. It accepts suffixes. (e.g. -c 50G). Default %s.
\t-u, --units\tAll  sizes are output in these units: (k)ilobytes, (m)egabytes,  (g)igabytes,  (t)erabytes, (p)etabytes. Capitalise  to use multiples of 1000 (S.I.) instead of 1024.
\t-p, --pools\tRestricts the pools to check to a list sparated by commas. (e.g. pool1,pool2)
\t-O, --VO\tRestricts the pools to check to a specific VO. By name or id. (e.g. dteam)

For each pool define with the \"--pool\" option, the probe will return the current used and free space. If the free space is under warning or critical threshold, the probe trigger an alert

Description of work executed by the probe:

\t1. Retrieve information on the monitored pool using the command "dpm_getpools()"
\t2. Parse the result of the command to find
\t\tThe free space
\t\tThe used space
\t3. Return values to nagios
\t\tWarning alert is triggered if free space is under 5 GBytes
\t\tCritical alert is triggered if free space is under 1 GBytes
""" % (DEFAULT_WARNING, DEFAULT_CRITICAL)

  # Methods
  def __init__(self, opt = {}, args = []):
    """
    Constructor

    @param opt  Contains a dictionary with the long option name as the key, and the argument as value
    @param args Contains the arguments not associated with any option
    """
    # Pools
    self.pools = None
    if "pools" in opt:
      self.pools = opt["pools"].split(',')

    # Get warning and critic options
    opt_warning  = self.DEFAULT_WARNING
    opt_critical = self.DEFAULT_CRITICAL
    opt_units = self.DEFAULT_UNITS

    if "warning" in opt:
      opt_warning = opt["warning"]
    if "critical" in opt:
      opt_critical = opt["critical"]
    if "units" in opt:
      opt_units = opt["units"]

    self.pool_warning = real_bytes(opt_warning)
    self.pool_critical = real_bytes(opt_critical)
    self.units = opt_units

    # VO
    self.vo_id = self.vo = None
    if "VO" in opt:
      try:
        self.vo_id = int(opt["VO"])
      except ValueError:
        # Passed as name
        self.vo = opt["VO"]

  # if the --pools option is set, restrain the results to that list 
  def filter_pools(self, pool_array):
    filtering_pool_array = []
    for pool in pool_array:
      if pool.poolname in self.pools:
        filtering_pool_array.append(pool)
    return filtering_pool_array

  def main(self):
    """
    Test code itself. May raise exceptions.

    @return A tuple (exit code, message, performance)
    """
    # Get pools
    (stat, pool_array) = dpm.dpm_getpools()
    if stat != 0:
      return (EX_UNKNOWN, "Could not get the pools (%d)" % dpm.C__dpm_errno(), None)

    # Get VO id or name (here because of a bug in dpm module)
    import dpm2
    if self.vo_id:
      self.vo = dpm2.dpns_getgrpbygid(self.vo_id)
    elif self.vo:
      self.vo_id = dpm2.dpns_getgrpbynam(self.vo)
    del dpm2
  
    if self.vo_id: 
      debug("VO: %s (%d)" % (self.vo, self.vo_id))

    # For each pool, get filesystems
    exit = EX_OK
    msg  = ""
    critical, warning = 0, 0
    performance = ""

    # if the --pools option is set, restrain the results to that list 
    if self.pools is not None:
      pool_array = self.filter_pools(pool_array)

    for pool in pool_array:
      if not self.vo_id or self.vo_id in pool.gids:
        if pool.free < self.pool_critical:
          exit = EX_CRITICAL
          msg += "%s in a critical state\n " % pool.poolname
          critical += 1
        elif pool.free < self.pool_warning:
          exit = safe_exit(EX_WARNING, exit)
          msg += "%s in a warning state\n " % pool.poolname
          warning += 1 

        performance += "'" + pool.poolname + "'="
        performance += str(standard_units(pool.free, self.units))
        if self.units is not None: 
          performance += self.units
        performance += "B;" + str(standard_units(self.pool_warning, self.units)) + ";"
        performance += str(standard_units(self.pool_critical, self.units)) + ";"
        performance += "0;" + str(standard_units(pool.capacity, self.units))

      # end for if not self._vo_id or self.vo_id in pool.groups
    # end for pool

    msg = "%d critical(s) ans %d warning(s) found\n%s" % (critical, warning, msg)

    if exit == EX_OK:
      msg = "All DPM pools are OK"
    else:
      msg = msg[:-2] # Supress last "; "

    if self.vo_id:
      msg += " (VO %s)" % self.vo

    return (exit, msg, performance)

# When called directly
if __name__ == "__main__":
  run(check_dpm_pool)

