#!/usr/bin/env python

def main():

    from subprocess import check_output
    from subprocess import CalledProcessError
    import shlex
    import json
    import os
    import re

    if os.path.exists('/etc/ceph'):
        grain = {}
        grain["ceph"] = {}
        conf_dir = '/etc/ceph/'
        for filename in os.listdir(conf_dir):
            if filename.endswith(".conf"):
                cluster_name = re.search('(.+?).conf', filename).group(1)
                break
        conf_file = conf_dir + cluster_name + '.conf'

        # get the fsid from config file, for salt-formulas to filter on in case of multiple ceph clusters
        with open(conf_file, 'r') as conf_fh:
            for line in conf_fh.read().splitlines():
                if 'fsid' in line:
                    attr = shlex.split(line)
                    grain['ceph']['fsid'] = attr[2]

        # crush id
        cmd = "ceph osd crush tree | grep `hostname -s` | awk '{print $1}'"
        crush_id = check_output(cmd, shell=True).strip()
        if crush_id:
            grain['ceph']['crush_id'] = crush_id

        # osd
        if os.path.exists('/var/lib/ceph/osd'):
            devices = {}
            for osd in os.listdir('/var/lib/ceph/osd'):

                try:
                    with open('/var/lib/ceph/osd/%s/whoami' % osd) as file:
                        dev_id = file.readline().strip()
                    with open('/var/lib/ceph/osd/%s/fsid' % osd) as file:
                        fsid = file.readline().strip()
                    path = os.readlink('/var/lib/ceph/osd/%s/block' % osd)
                except IOError:
                    continue            

                cmd = "lsblk -no pkname %s" % path
                dev = check_output(cmd, shell=True).strip()
                if dev:
                    dev = '/dev/%s' % dev
                else:
                    cmd = "pvs  | grep %s | awk '{print $1}'" % path.split('/')[-2]
                    dev = check_output(cmd, shell=True).strip()
                devices[dev_id] = {}
                devices[dev_id]['fsid'] = fsid
                devices[dev_id]['path'] = path
                devices[dev_id]['dev'] = dev
                
            cmd = "ceph osd tree --format json"
            osd_tree_output = check_output(cmd, shell=True).decode("utf-8")
            osd_tree_output = json.loads(osd_tree_output)
            for osd in osd_tree_output['nodes']:
                dev_id = str(osd['id'])
                if dev_id >= 0 and devices.has_key(dev_id):
                    devices[dev_id]['weight'] = osd['crush_weight']
                    devices[dev_id]['class'] = osd['device_class']

            if devices:
                grain['ceph']['ceph_disk'] = devices

        # keyrings
        directory = '/etc/ceph/'
        keyrings = {}
        if os.path.isdir(directory):
            for filename in os.listdir(directory):
                if filename.endswith(".keyring") and re.search(".client.", filename):
                    keyring_output = open(os.path.join(directory, filename), "r")
                    keyring_name = re.search('(.+?).client.(.+?).keyring', filename).group(2)
                    if keyring_output:
                        keyrings[keyring_name] = {}
                        for line in keyring_output:
                            attr = shlex.split(line)
                            if attr:
                                if attr[0] == 'key':
                                    keyrings[keyring_name]['key'] = attr[2]
                                if attr[0] == 'caps' and 'caps' in keyrings[keyring_name]:
                                    keyrings[keyring_name]['caps'][attr[1]] = attr[3]
                                elif attr[0] == 'caps' and 'caps' not in keyrings[keyring_name]:
                                    keyrings[keyring_name]['caps'] = {}
                                    keyrings[keyring_name]['caps'][attr[1]] = attr[3]
            if keyrings:
                grain["ceph"]["ceph_keyring"] = keyrings

        # mon keyring
        hostname = check_output("hostname", shell=True).rstrip()
        filepath = "/var/lib/ceph/mon/{0}-{1}/keyring".format(cluster_name, hostname)
        if os.path.isfile(filepath):
            mon_key_output = open(filepath, "r")
            if mon_key_output:
                keyrings['mon'] = {}
                for line in mon_key_output:
                    attr = shlex.split(line)
                    if attr:
                        if attr[0] == 'key':
                            keyrings['mon']['key'] = attr[2]
                        if attr[0] == 'caps' and 'caps' in keyrings['mon']:
                            keyrings['mon']['caps'][attr[1]] = attr[3]
                        elif attr[0] == 'caps' and 'caps' not in keyrings['mon']:
                            keyrings['mon']['caps'] = {}
                            keyrings['mon']['caps'][attr[1]] = attr[3]
                grain["ceph"]["ceph_keyring"] = keyrings

        return grain
    else:
        return None