#!/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"

  # 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=",
                         "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-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)

The accepted suffixes are K(ibibyte), M(ebibyte), G(ibibyte), T(ebibyte) and P(ebibyte).

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
    if "warning" in opt:
      opt_warning = opt["warning"]
    if "critical" in opt:
      opt_critical = opt["critical"]

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

    # 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"]

  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  = ""
    performance = []

    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 pool free space is under the critical limit; " % pool.poolname
        elif pool.free < self.pool_warning:
          exit = safe_exit(EX_WARNING, exit)
          msg += "%s pool free space is under the warning limit; " % pool.poolname

        performance.append("'%s'=%dB;%d:;%d:;0;%d" % (pool.poolname, pool.free, self.pool_warning, self.pool_critical, pool.capacity))
      
      # end for if not self._vo_id or self.vo_id in pool.groups
    # end for pool

    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

    performance = " ".join(performance)

    return (exit, msg, performance)

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

