#!/usr/bin/env python
import paramiko
import sys

# LSHW parsing
def ls_parse_branch(lines, depth=1):
    data = {}
    while len(lines):
        line = lines.pop(0)
        if line[depth:].startswith('   ') and ':' in line:
            index = line.index(':')
            key = line[depth+3:index]
            value = line[index+2:]
            data[key] = value
        elif line[depth:].startswith(' *-'):
            key = line[depth+3:]
            # Check if key has a status
            status = None
            if ' ' in key:
                key, status = key.split(' ')

            # Check if key is a list member
            index = None
            if ':' in key:
                key, index = key.split(':')
                index = int(index)

            lines, sub_branch = ls_parse_branch(lines, depth+3)
            if status is not None:
                sub_branch['status'] = status
            if index is not None:
                if key not in data:
                    data[key] = []
                data[key].insert(index, sub_branch)
            else:
                data[key] = sub_branch
        else:
            return [line] + lines, data
    return [], data


def lshw_parse(lines):
    # If string given, split into list of lines
    if isinstance(lines, str):
        lines = lines.splitlines()

    # First line is hostname
    hostname = lines.pop(0).strip()

    # Parse remaining lines
    data = ls_parse_branch(lines)[1]

    data['hostname'] = hostname
    return data

# LSPCI parsing
def lspci_tree(lines):
    data = {}
    parameters = ['LnkCap', 'LnkSta', 'Kernel modules']
    while len(lines):
        line = lines.pop(0)
        if not line.startswith('\t'):
            return [line] + lines, data
        else:
            for parameter in parameters:
                if line.startswith('\t\t{}:'.format(parameter)):
                    data[parameter] = line[len(parameter)+3:]
                elif line.startswith('\t{}:'.format(parameter)):
                    data[parameter] = line[len(parameter)+2:]
    return lines, data

def lspci_device(lines):
    line = lines.pop(0)
    bus_separator = line.index(' ')
    desc_seperator = bus_separator + line[bus_separator:].index(':')
    bus = line[:bus_separator]
    lines, data = lspci_tree(lines)
    data['Type'] = line[bus_separator+1:desc_seperator]
    data['Description'] = line[desc_seperator+2:]
    return lines, bus, data

def lspci_parse(lines):
    devices = {}
    # Ensure lines is a list of non-empty strings
    if isinstance(lines, str):
        lines = [x for x in lines.splitlines() if x.strip()]
    else:
        lines = [x for x in lines if x.strip()]

    while len(lines):
        lines, bus, data = lspci_device(lines)
        if not data['Description'].startswith('IBM Device 0'):
            devices[bus] = data
    return devices



class G2_SSH(object):
    ''' SSH-based interface to OpenBMC system '''
    def __init__(self, address, username='root', password='0penBmc',
                 sshtimeout=10):
        self.address = address
        self.ssh = paramiko.SSHClient()
        self.ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        self.ssh.connect(address, timeout=sshtimeout)

    def command(self, command):
        ''' Runs a shell command, returns (stdin, stdout, stderr) handles '''
        return self.ssh.exec_command(command)

    def call(self, command):
        stdin, stdout, stderr = self.command(command)
        return stdout.read(), stderr.read()

    def host_state(self):
        ''' Returns current state of host '''
        stdin, stdout, stderr = self.command('/usr/sbin/obmcutil state')
        return stdout.read().rstrip().split(' = ')[-1]

    def uptime(self):
        ''' Returns BMC uptime '''
        stdin, stdout, stderr = self.command('uptime')
        return stdout.read().strip()

    def meminfo(self):
        ''' Get BMC memoey information '''
        stdin, stdout, stderr = self.command('cat /proc/meminfo')
        info = {}
        for line in stdout.readlines():
            key, value = line.split(':')
            info[key] = value.strip()
        return info

    def lshw(self):
        stdout, stderr = self.call('lshw')
        return lshw_parse(stdout)

    def lspci(self):
        stdout, stderr = self.call('lspci -vvv')
        return lspci_parse(stdout)

try:
    target = sys.argv[1]
except:
    print 'Usage: {} <target host>'.format(sys.argv[0])
    sys.exit(1)

print "Connecting to {}...".format(target)
boxen = G2_SSH(target)
print 'Getting info from lshw...'
hw = boxen.lshw()
print 'Getting info from lspci...'
pci = boxen.lspci()

print '\nProc Info:'
core_counts = {}
cores_total = 0
for core in hw['core']['cpu']:
    if 'slot' in core:
        slot = core['slot'][-1]
        if slot not in core_counts:
            core_counts[slot] = 0
        core_counts[slot] += 1
    cores_total += 1

for socket in sorted(core_counts.keys()):
    print ' Cores in socket {}: {}'.format(socket, core_counts[socket])
print ' {} Cores total'.format(cores_total)

print '\nMemory Info:'
if 'bank' in hw['core']['memory']:
    print ' {} RAM from {} DIMMs'.format(hw['core']['memory']['size'],
                                       len(hw['core']['memory']['bank']))
else:
    print ' {} RAM total'.format(hw['core']['memory'])

print '\nPCI Devices:'
for bus in pci:
    device = pci[bus]
    print ' {} - {}'.format(device['Type'], device['Description'])
    print '  PCI Bus: {}'.format(bus)
    if 'LnkCap' in device:
        cap = ', '.join([x.strip() for x in device['LnkCap'].split(',')[1:3]])
        print '  Link Capability: {}'.format(cap)
    if 'LnkSta' in device:
        sta = ', '.join([x.strip() for x in device['LnkSta'].split(',')[0:2]])
        print '  Link Status: {}'.format(sta)
    if 'Kernel modules' in device:
        print '  Kernel Modules: {}'.format(device['Kernel modules'])
    print ''