#!/usr/local/bin/python
# Copyright (c) 2012, Jakob Borg
# All rights reserved.
#
# 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.
#    * The name of the author may not be used to endorse or promote products
#      derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY JAKOB BORG ''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 JAKOB BORG 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.
# Initial code taken from: https://github.com/jm66/solaris-extra-snmp

import sys
import commands
import re
import snmpresponse
import json
import socket
import libzfs


sys.path.append("/usr/local/www")

from freenasUI.tools.arc_summary import get_Kstat, get_arc_efficiency

BASE_OID = '.1.3.6.1.4.1.25359.1'
ARC = get_arc_efficiency(get_Kstat())
FREENASSNMPDSOCK = '/var/run/freenas-snmpd.sock'

size_dict = {
    "K": 1024,
    "M": 1048576,
    "G": 1073741824,
    "T": 1099511627776
}


def get_from_freenas_snmpd_sock(val_to_obtain):
    data = b''
    try:
        s = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        s.connect(FREENASSNMPDSOCK)
        s.sendall(val_to_obtain)
        while True:
            text = s.recv(4096)
            if text == b'':
                break
            else:
                data += text
    except socket.error:
        pass
    finally:
        s.close()
    try:
        data = json.loads(data.decode('utf8'))
    except (ValueError, UnicodeDecodeError):
        data = {}
    return data


FSNMPDATA = get_from_freenas_snmpd_sock("get_all")


class ArgumentValidationError(ValueError):
    """
    Raised when the type of an argument to a function is not what it should be.
    """

    def __init__(self, arg_num, func_name, accepted_arg_type):
        self.error = 'The {0} argument of {1}() is not a {2}'.format(
            arg_num, func_name, accepted_arg_type)

    def __str__(self):
        return self.error


def unprettyprint(ster):
    """
    Method to convert 1K --> 1024 and so on...
    """
    num = 0.0
    try:
        num = float(ster)
    except:
        try:
            num = float(ster[:-1]) * size_dict[ster[-1]]
        except:
            pass
    return long(num)


def kstat(name):
    output = commands.getoutput("sysctl kstat." + name)
    try:
        return int(re.split("\s+", output)[1])
    except:
        return 0


def zfs_arc_size(oid):
    # KB
    return ('gauge', kstat("zfs.misc.arcstats.size") / 1024)


def zfs_arc_data(oid):
    # KB
    return ('gauge', kstat("zfs.misc.arcstats.data_size") / 1024)


def zfs_arc_meta(oid):
    # KB
    return ('gauge', kstat("zfs.misc.arcstats.arc_meta_used") / 1024)


def zfs_arc_hits(oid):
    # 32 bit counter
    return ('counter', kstat("zfs.misc.arcstats.hits") % 2**32)


def zfs_arc_misses(oid):
    # 32 bit counter
    return ('counter', kstat("zfs.misc.arcstats.misses") % 2**32)


def zfs_arc_miss_percent(oid):
    # percentage (floating point precision wrapped as a string)
    arc_hits = kstat("zfs.misc.arcstats.hits")
    arc_misses = kstat("zfs.misc.arcstats.misses")
    arc_read = arc_hits + arc_misses
    if (arc_read > 0):
        hit_percent = float(100 * arc_hits / arc_read)
        miss_percent = 100 - hit_percent
        return ('string', str(miss_percent))
    return ('string', "0")


def zfs_arc_c(oid):
    # KB
    return ('gauge', kstat("zfs.misc.arcstats.c") / 1024)


def zfs_arc_p(oid):
    # KB
    return ('gauge', kstat("zfs.misc.arcstats.p") / 1024)


def zfs_arc_cache_hit_ratio(oid):
    # percentage (floating point precision wrapped as a string)
    return ('string', ARC['cache_hit_ratio']['per'][:-1])


def zfs_arc_cache_miss_ratio(oid):
    # percentage (floating point precision wrapped as a string)
    return ('string', ARC['cache_miss_ratio']['per'][:-1])


def zilstatd_ops(interval):
    res = long(0)
    try:
        res = FSNMPDATA["zil_data"][str(interval)]['ops']
    except KeyError:
        pass
    return res


# Note: Currently only 1 second interval and "all" is supported
# to add more make the appropriate worker in gui/tools/freenas-snmpd.py
def zpoolio_interval(interval):
    res = {}
    try:
        res = FSNMPDATA["zpool_data"][str(interval)]
    except KeyError:
        pass
    return res


def zfs_zilstat_ops1(oid):
    return ('counter64', zilstatd_ops(1))


def zfs_zilstat_ops5(oid):
    return ('counter64', zilstatd_ops(5))


def zfs_zilstat_ops10(oid):
    return ('counter64', zilstatd_ops(10))


def zfs_l2arc_size(oid):
    # KB
    return ('gauge', kstat("zfs.misc.arcstats.l2_size") / 1024)


def zfs_l2arc_hits(oid):
    # 32 bit counter
    return ('counter',
            kstat("zfs.misc.arcstats.l2_hits") % 2**32)


def zfs_l2arc_misses(oid):
    # 32 bit counter
    return ('counter',
            kstat("zfs.misc.arcstats.l2_misses") % 2**32)


def zfs_l2arc_write(oid):
    # 32 bit KB counter
    return ('counter',
            kstat("zfs.misc.arcstats.l2_write_bytes") / 1024 % 2**32)


def zfs_l2arc_read(oid):
    # 32 bit KB counter
    return ('counter',
            kstat("zfs.misc.arcstats.l2_read_bytes") / 1024 % 2**32)


def zfs_segregate(zfs_dataset):
    """
    A function to obtain and segregte all the datsets (children) and zvols
    in the provided (`zfs_dataset`) dataset. The best example to use this
    is to provide it with a zfs pool's root dataset and it will return a
    a tuple of (datsets, zvols) where each of them is a list.

    Note: Please make sure that the input to this function (`zfs_dataset`)
    is of type: libzfs.ZFSDataset
    """
    if type(zfs_dataset) is not libzfs.ZFSDataset:
        raise ArgumentValidationError(1, 'zfs_segregate', libzfs.ZFSDataset)
    zvols = []
    datasets = []
    if zfs_dataset.properties['type'].value == 'volume':
        # since zvols do not have children lets just return now
        return datasets, [zfs_dataset]
    else:
        datasets = [zfs_dataset]
    for x, y in map(zfs_segregate, list(zfs_dataset.children)):
        datasets.extend(x)
        zvols.extend(y)
    return datasets, zvols


# nymnetworks     OBJECT IDENTIFIER ::= {enterprises 25359}
# zfs             OBJECT IDENTIFIER ::= {nymnetworks 1}
# fs              OBJECT IDENTIFIER ::= {zfs 1}
# arc             OBJECT IDENTIFIER ::= {zfs 2}
# l2arc           OBJECT IDENTIFIER ::= {zfs 3}
# io              OBJECT IDENTIFIER ::= {zfs 4}
# vols            OBJECT IDENTIFIER ::= {zfs 5}
# zil             OBJECT IDENTIFIER ::= {zfs 6}
# ds              OBJECT IDENTIFIER ::= {zfs 7}


result = [
    (BASE_OID + '.2.1.0', zfs_arc_size),
    (BASE_OID + '.2.2.0', zfs_arc_meta),
    (BASE_OID + '.2.3.0', zfs_arc_data),
    (BASE_OID + '.2.4.0', zfs_arc_hits),
    (BASE_OID + '.2.5.0', zfs_arc_misses),
    (BASE_OID + '.2.6.0', zfs_arc_c),
    (BASE_OID + '.2.7.0', zfs_arc_p),
    (BASE_OID + '.2.8.0', zfs_arc_miss_percent),
    (BASE_OID + '.2.9.0', zfs_arc_cache_hit_ratio),
    (BASE_OID + '.2.10.0', zfs_arc_cache_miss_ratio),

    (BASE_OID + '.3.1.0', zfs_l2arc_hits),
    (BASE_OID + '.3.2.0', zfs_l2arc_misses),
    (BASE_OID + '.3.3.0', zfs_l2arc_read),
    (BASE_OID + '.3.4.0', zfs_l2arc_write),
    (BASE_OID + '.3.5.0', zfs_l2arc_size),

    (BASE_OID + '.6.1.0', zfs_zilstat_ops1),
    (BASE_OID + '.6.2.0', zfs_zilstat_ops5),
    (BASE_OID + '.6.3.0', zfs_zilstat_ops10),
]

all_iostat = zpoolio_interval("all")
onesec_iostat = zpoolio_interval(1)
zfs = libzfs.ZFS()
pools = [pool for pool in zfs.pools]
datasets = []
zvols = []
for res in map(zfs_segregate, [pool.root_dataset for pool in pools]):
    # Excluding the first item in the datasets list as its always the root_dataset
    # and we already have the stats on that (i.e. the pool!)
    datasets.extend(res[0][1:])
    zvols.extend(res[1])

for i, zpool in enumerate(pools):
    stri = str(i)
    pool = zpool.name
    pool_health = zpool.properties['health'].value
    # Dividing by 1024 to ge to KB
    pool_used = unprettyprint(zpool.root_dataset.properties['used'].value) / 1024
    pool_available = unprettyprint(zpool.root_dataset.properties['available'].value) / 1024
    pool_size = unprettyprint(zpool.properties['size'].value) / 1024
    result.extend([
        (BASE_OID + '.1.1.' + stri, ('string', pool)),
        (BASE_OID + '.1.2.' + stri, ('gauge', pool_available)),
        (BASE_OID + '.1.3.' + stri, ('gauge', pool_used)),
        (BASE_OID + '.1.4.' + stri, ('string', pool_health)),
        (BASE_OID + '.1.5.' + stri, ('gauge', pool_size)),
        (BASE_OID + '.1.12.' + stri, ('gauge', pool_available / 1024)),
        (BASE_OID + '.1.13.' + stri, ('gauge', pool_used / 1024)),
        (BASE_OID + '.1.14.' + stri, ('gauge', pool_size / 1024)),
        (BASE_OID + '.1.15.' + stri, lambda oid, fs=pool: ('counter64', all_iostat[fs]['opread'])),
        (
            BASE_OID + '.1.16.' + stri,
            lambda oid, fs=pool: ('counter64', all_iostat[fs]['opwrite'])
        ),
        (BASE_OID + '.1.17.' + stri, lambda oid, fs=pool: ('counter64', all_iostat[fs]['bwread'])),
        (BASE_OID + '.1.18.' + stri, lambda oid, fs=pool: ('counter64', all_iostat[fs]['bwrite'])),
        (
            BASE_OID + '.1.19.' + stri,
            lambda oid, fs=pool: ('counter64', onesec_iostat[fs].get('opread', 0))
        ),
        (
            BASE_OID + '.1.20.' + stri,
            lambda oid, fs=pool: ('counter64', onesec_iostat[fs].get('opwrite', 0))
        ),
        (
            BASE_OID + '.1.21.' + stri,
            lambda oid, fs=pool: ('counter64', onesec_iostat[fs].get('bwread', 0))
        ),
        (
            BASE_OID + '.1.22.' + stri,
            lambda oid, fs=pool: ('counter64', onesec_iostat[fs].get('bwrite', 0))
        )
    ])

for i, zvol in enumerate(zvols):
    stri = str(i)
    volsize = unprettyprint(zvol.properties['volsize'].value) / 1024
    vol_used = unprettyprint(zvol.properties['used'].value) / 1024
    vol_available = unprettyprint(zvol.properties['available'].value) / 1024
    result.extend([
        (BASE_OID + '.5.1.' + stri, ('string', zvol.name)),
        (BASE_OID + '.5.2.' + stri, ('gauge', vol_available)),
        (BASE_OID + '.5.3.' + stri, ('gauge', vol_used)),
        (BASE_OID + '.5.4.' + stri, ('gauge', volsize)),
        (BASE_OID + '.5.12.' + stri, ('gauge', vol_available / 1024)),
        (BASE_OID + '.5.13.' + stri, ('gauge', vol_used / 1024)),
        (BASE_OID + '.5.14.' + stri, ('gauge', volsize / 1024))
    ])

for i, ds in enumerate(datasets):
    stri = str(i)
    ds_used = unprettyprint(ds.properties['used'].value) / 1024
    ds_available = unprettyprint(ds.properties['available'].value) / 1024
    ds_size = ds_used + ds_available
    result.extend([
        (BASE_OID + '.7.1.' + stri, ('string', ds.name)),
        (BASE_OID + '.7.2.' + stri, ('gauge', ds_available)),
        (BASE_OID + '.7.3.' + stri, ('gauge', ds_used)),
        (BASE_OID + '.7.4.' + stri, ('gauge', ds_size)),
        (BASE_OID + '.7.12.' + stri, ('gauge', ds_available / 1024)),
        (BASE_OID + '.7.13.' + stri, ('gauge', ds_used / 1024)),
        (BASE_OID + '.7.14.' + stri, ('gauge', ds_size / 1024))
    ])


operation = sys.argv[1]
req_oid = sys.argv[2]

snmpresponse.respond_to(operation, req_oid, result)
