#!/usr/bin/env python3
#
#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation; either version 2 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License along
#   with this program; if not, write to the Free Software Foundation, Inc.,
#   51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

# Marc Schoechlin <ms@256bit.org>

import argparse
import glob
import json
import os
import re
import sys

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

config_default = {
    "regex_includes_name": [".*"],
    "regex_includes_model": [".*"],
    "regex_excludes_name": ["loop.*", "fd.*", "sr.*", "dm.*", "ram.*" ],
    "regex_excludes_model": [],
    "smartctl_subdevices": {},
}


####################################################################################
#### HELPERS

def check_match(matchfor, device_name, regexes):
    ret = False
    for regex in regexes:
        regex = regex.strip()
        if re.match(r"^%s$" % regex, device_name):
            if args.debug:
                sys.stderr.write("%s match: %s with %s\n" % (matchfor, device_name, regex))
            ret = True
    return ret


def read_config(config_file):
    config = config_default
    if os.path.exists(config_file):
        if args.debug:
            sys.stderr.write("reading config file %s\n" % config_file)
        try:
            fh = open(config_file, 'r')
            content = fh.read()

            # Remove multiline comments
            content = re.sub('^\s*/\*.*?\*/\s*', '', content, 0, re.DOTALL | re.MULTILINE)
            # Remove single line comments
            content = re.sub('^\s*//.*$\n', '', content, 0, re.MULTILINE)

            config = json.loads(content)

            fh.close()
        except FileNotFoundError:
            sys.stderr.write("Config file does not exist\n")
        except PermissionError:
            sys.stderr.write("Could not read config file\n")
        except json.decoder.JSONDecodeError:
            sys.stderr.write("Could not parse JSON config file\n")
        except:
            sys.stderr.write("An unexpected error occured while reading config file\n")

    config.setdefault("smartctl_subdevices", {})
    return config


def is_real_hardware_device(device):
    if os.path.exists("/sys/block/%s/device/vendor" % device):
        return True
    else:
        return False

####################################################################################
#### MAIN


parser = argparse.ArgumentParser(
    description='perform device disovery for zabbix',
    formatter_class=argparse.RawTextHelpFormatter,
    epilog='''Example json configuration file:

i.e. /etc/zabbix/item_zabbix_device_discovery.json:
---
{
  "regex_includes_name": [
    ".*"
  ],
  "regex_includes_model": [
    ".*"
  ],
  "regex_excludes_name": [
    "loop.*", "fd.*", "sr.*" "md.*"
  ],
  "regex_excludes_model": [
  ],
  "smartctl_subdevices": {
    "sda": [
      "-d megaraid,0",
      "-d megaraid,1"
    ]
  }
}
---

Options are evaluated like this:
- stage 1: include the following devices by regex
- stage 2: exclude the following devices from step 1 by regex

    '''
)
parser.add_argument(
    '--debug',
    help='Output debug information',
    action='store_true',
)

parser.add_argument(
    '--smart',
    help='Discovery for smartctl',
    action='store_true',
)


parser.add_argument(
    '--only_rawdisk',
    help='Output only hardware devices',
    action='store_true',
)


parser.add_argument(
   '--config',
   nargs='?',
   type=str,
   help='configuration file',
   default=None,
   required=False,
)

args = parser.parse_args()

devices_discovered = {'data': []}

# read the config
if args.config is not None:
    config = read_config(args.config)
else:
    config = config_default

if args.debug:
    sys.stderr.write("** Effective config:\n%s\n" % json.dumps(config, indent=2))



# perform discovery
for filename in glob.glob('/sys/block/*'):
    if not os.path.islink(filename):
        continue
    device_name = os.path.basename(filename)

    device_model_file = "/sys/block/%s/device/model" % device_name
    device_model = "UNDEFINED"
    
    if os.path.isfile(device_model_file):
        with open(device_model_file) as f_model:
                device_model = f_model.readline()

    device_size_file = "/sys/block/%s/size" % device_name 
    device_size = 0
    if os.path.isfile(device_size_file):
        with open(device_size_file) as f_size:
               device_size = int(f_size.readline().strip())
    if device_size == 0:
        if args.debug:
             sys.stderr.write("ignoring device %s, its size is 0\n" % device_name)
        continue

    if args.debug:
        sys.stderr.write("** DEVICE %s / MODEL %s\n" % (device_name,device_model))

    # Check matched devices
    if not check_match("include", device_name, config["regex_includes_name"]):
        continue

    if not check_match("include", device_model, config["regex_includes_model"]):
        continue

    # Check excluded devices
    if check_match("exclude", device_name, config["regex_excludes_name"]):
        continue

    if check_match("exclude", device_model, config["regex_excludes_model"]):
        continue

    if args.only_rawdisk and not is_real_hardware_device(device_name):
        continue

    if args.smart and device_name in config["smartctl_subdevices"]:
        for smartarg in config["smartctl_subdevices"][device_name]:
            device = {}
            device['{#BLOCKDEVICE}'] = device_name +  " " + smartarg
            devices_discovered['data'].append(device)
    else:
        device = {}
        device['{#BLOCKDEVICE}'] = device_name
        devices_discovered['data'].append(device)



if args.debug:
    print(json.dumps(devices_discovered, sort_keys=True, indent=2))
else:
    print(json.dumps(devices_discovered, sort_keys=True))

# vim: ai et ts=2 shiftwidth=2 expandtab
