#
# This script is designed to help identify physical drives in use in zpools.
# Generally this drive info includes serial number as printed on drive.
#
# Executing this script does not execute any commands that change the
# stored on disk state of the operating system and is considered low risk
# to run. But a precuation is to not run any script you find on the internet
# or are given unless you understand what it does first.
#
# This script dumps out a zpool with its related vdevs and drive information.
# It uses the following commands to collect the information it displays.
#   "zpool status"
#   "camcontrol devlist"
#   "glabel status"
#   "smartctl -i /dev/<device>"  - this is run for each device to lookup serial
#
# To use it copy this script to the server login as root and execute:
#   "python zpool-drive-info.py <poolname>"
#
# Where <poolname> is the name of the pool you want information on.
#
# Python is something I only dabble in and havnt touched in many many years
# so if this python source offends people I am sorry, but I really didnt want
# to write this Bourne shell or Awk.
#
# Copyright Robin Luiten 2016.
# This script is free to use and distribute.
# If you fix anything I would appreciate being notified.
#
#   Written on FreeNAS 8.30 running on HP Microserver N36L.
#   Tested on FreeNAS 8.30 running on HP Microserver N36L.
#   Tested on FreeNAS 9.3.
#   Tested on FreeNAS 9.1.
#
# 2016/03/31 Created v1.0.0
# 2016/04/02 v1.0.1
#   Fixed bug in lookup of device id.
# 2016/04/02 v1.0.3
#   Now looks up serial number of drive via smartctl
# 2016/04/02 v1.0.4
#   Removed old output code
# 2016/04/02 v1.0.5
#   Tweak output format, udated FreeNAS versions it was tested on.
# 2016/04/03 v1.0.6
#   Fixed comment.
# 2016/04/17 v1.0.6
#   Fixed incompatibilty with result of `camcontrol devlist`
# 2016/04/18 v1.0.7
#   Bug with single vdev with single disk. They dont have a vdev1 basically so
#   it doesnt find a vdev2 level 2.
#
import os
import re
import subprocess
import sys
# import test

if len(sys.argv) != 2:
    print 'USAGE: zfsls.py poolname'
    exit(1)

dev1Counter=0
tank = sys.argv[1]
nameColumn = -1
singleVdevDiskMode = False
zpool = subprocess.check_output(["zpool", "status", tank])
camcontrol = subprocess.check_output(["camcontrol", "devlist"])
glabel = subprocess.check_output(["glabel", "status"])

# zpool = test.test['finch']['zpooltank1']
# zpool = test.zpool0;
# camcontrol = test.test['finch']['camcontrol']
# camcontrol = test.camcontrol1;
# glabel = test.test['finch']['glabel']
# glabel = test.glabel0;

zpoolStructure = {}
for line in zpool.splitlines():
    # print "line: \"" + line + "\""
    if nameColumn < 0:
        column = line.find("NAME")
        if column >= 0:
            nameColumn = column
            # print "NAME column", nameColumn
    else:
        if len(line) > nameColumn:
            leftTrimLine = line[nameColumn:]
            match = re.search(r'^\s*([\S]*)', leftTrimLine)
            if match and match.start(1) == 0:  # found poolname
                poolName = match.group(1)
                dev1 = ""
                dev2 = ""

            # found dev level 1, uniquify its name with counter
            if match and match.start(1) == 2:  # found dev level 1
                newDev1 = match.group(1)
                dev1Counter = dev1Counter + 1
                dev1 = newDev1 + " (" + str(dev1Counter) + ")"
                dev2 = ""
                singleVdevDiskMode = newDev1.find("gptid") >= 0

            if (match and match.start(1) == 4) or singleVdevDiskMode: # found dev level 2
                dev2 = match.group(1)

            if poolName:
                if not zpoolStructure.has_key(poolName):
                    zpoolStructure[poolName] = {}
                if dev1:
                    if not zpoolStructure[poolName].has_key(dev1):
                        zpoolStructure[poolName][dev1] = {}
                    if dev2:
                        if not zpoolStructure[poolName][dev1].has_key(dev2):
                            zpoolStructure[poolName][dev1][dev2] = True
        else:
            break
# print "zpoolStructure", zpoolStructure

# print camcontrol
camcontrolStucture = {}
for line in camcontrol.splitlines():
    # print "a", "\"" + line + "\""
    match = re.search(r'^\s*(.*)\s*\sat\s.*\(([^,]+),([^)]+)\)$', line)
    if match: # group 2 and 3 contain values inside () at end of line.
        # just add both to table, the one that is status "pass*" and not device
        # will just wont be fond in device lookup later.
        camcontrolStucture[match.group(2)] = match.group(1).strip()
        camcontrolStucture[match.group(3)] = match.group(1).strip()
# print "camcontrolStucture", camcontrolStucture

# print "glabel"
glabelStucture = {}
for line in glabel.splitlines():
    # print "a", "\"" + line + "\""
    match = re.search(r'^\s*(\S+).*\s(\S+)\s*$', line)
    if match:
        justDevice = re.sub(r'p[\d]+$',r'', match.group(2))
        # print match.group(2), justDevice
        glabelStucture[match.group(1)] = justDevice
# print "glabelStucture", glabelStucture

def getDeviceInfo(deviceId):
    keys = camcontrolStucture.keys()
    for key in keys:
        if deviceId == key:
            return camcontrolStucture[key]

    return "Weird \"" + deviceId + "\" not found in device list"

def getSerial(deviceId):
    deviceId = deviceId
    serial = "Can't find serial for device /dev/" + deviceId
    info = subprocess.check_output(["smartctl", "-i", "/dev/" + deviceId]).splitlines()
    #info = "abced"
    for i in info:
        if i.find("Serial Number:") >= 0:
            serial = re.sub(r'^.*\s(\S+)$', r'\1', i)
    return serial

poolKeys = zpoolStructure.keys()
for poolName in sorted(poolKeys):
    print "      pool:  " + poolName
    poolStructure = zpoolStructure[poolName]
    dev1Keys = poolStructure.keys()
    for dev1Name in sorted(dev1Keys):
        print "      vdev:    " + dev1Name
        dev1Structure = zpoolStructure[poolName][dev1Name]
        dev2Keys = dev1Structure.keys()
        for dev2Name in sorted(dev2Keys):
            print
            print "        id:      " + dev2Name
            print "    device:        " + glabelStucture[dev2Name]
            print " driveinfo:        " + getDeviceInfo(glabelStucture[dev2Name])
            print "    serial:        " + getSerial(glabelStucture[dev2Name])
print