#!/usr/bin/python
#
# A NG lspci utility with some goodies
# Version 1.0
# Vaithiyanathan Sundaram
# Copyright (c) 2016 by Cisco Systems, Inc.
#
######################################################

import re
import subprocess
import sys
from collections import defaultdict
from optparse import OptionParser
import atexit
import time
import datetime

epilog = "Welcome to lspci-ng. This is created to recreate a wrapper around " \
         "lspci. lspci provides an unstructured " \
         "output of PCI devices, while one can use tree mode or can use -xx " \
         "or -xxx only when they can read bits " \
         "well. Personally, I am more comfortable with tabular form to " \
         "exactly give what we want.\n\n" \
         "In addition to the plain view mode,\n\n " \
         "raw mode: (-r) which will display in line by line manner useful " \
         "for grep. \n" \
         "inject mode: (-i) that can inject any form of AER in any order\n" \
         "verbose mode: (-v) can be switched on at any point of time which " \
         "will enable dmesg piping and gracefully " \
         "shuts it down\n" \
         "You can use -s and -d as you would use in the same way as lspci, " \
         "except wildcards.\n" \
         "Learn more about queries using -l option to list them. \n" \
         "If at all you find bugs or beautification, contact vaitsund"

parser = OptionParser(version='v1.0', epilog=epilog)
parser.add_option("-q", "--query", dest="query", help="Query mode")
parser.add_option("-l", "--query_list", dest="query_list", action='store_true',
                  default=False, help="Query list")
parser.add_option("-d", "--dev", dest="vendev",
                  help="Device id and vendor id, in lspci format venid:devid")
parser.add_option("-r", "--raw", dest="raw", action='store_true',
                  help="Raw mode without tabulation", default=False)
parser.add_option("-p", "--peek", dest="peek",
                  help="Raw mode without tabulation", default=None)
parser.add_option("-s", "--bdf", dest="bdf", help="Bus device function",
                  default=None)
parser.add_option("-i", "--inject", dest="aer_inject",
                  help="AER inject,\n"
                "example -i [LIST]\n"
                "lspci-ng -i <bdf>-<aer_type>-<aer_subtype>, ...\n"
                "You can add as many items as possible with comma\n"
                "The format for per inject is delimited by hyphen\n"
                "lspci-ng -i 42:00.0-UNCOR-MALF_TLP,00:03.0-UNCOR-MALF_TLP,...")
parser.add_option("-v", "--verbose", dest="verbose", action='store_true',
                  help="Enable dmesg piping, will only display logs on console, "
                       "aux won't work",
                  default=False)

options, args = parser.parse_args()

if options.raw and not options.query:
    print "Need query for raw mode use -q <query> or use -l for list " \
          "of queries along with -d venid:devid. You " \
          "may use -s for status along with -d <>"
    sys.exit(1)


def get_timestamp(line):
    try:
        return re.search('\[\s*(\d+.\d+)\]', line).groups()[0]
    except:
        return ''


if options.verbose:
    def disable_verbose():
        time.sleep(1)
        sys.stdout.flush()
        sys.stderr.flush()
        subprocess.call('dmesg -n 7', shell=True)
        subprocess.call('dmesg -D', shell=True)


    def enable_verbose():
        subprocess.call('dmesg -n 8', shell=True)
        subprocess.call('dmesg -E', shell=True)


    atexit.register(disable_verbose)  # Ensure we disable this, cleaner
    enable_verbose()

if options.aer_inject:

    def aer_inject():

        aer_types = ['UNCOR', 'COR']

        aer_subtypes = ['DLP', 'SDES', 'POISON_TLP', 'FCP', 'COMP_TIME',
                        'COMP_ABORT', 'UNX_COMP', 'RX_OVER',
                        'MALF_TLP', 'ECRC', 'UNSUP', 'ACS_VIOL']

        # Parse aer_inject
        injects = options.aer_inject.split(',')
        inject_list = [tuple(i.split('-')) for i in injects if
                       len(i.split('-')) == 3]
        print "Processing for %s" % inject_list

        file_template = ''
        for each_bdf, aer_type, aer_subtype in inject_list:
            if not aer_type in aer_types:
                print "Incorrect AER type"
                print '\n'.join(aer_types)
                sys.exit(1)

            if not aer_subtype in aer_subtypes:
                print "Incorrect AER subtype"
                print '\n'.join(aer_subtypes)
                sys.exit(1)

            # First step is to write the file
            file_template += 'AER\n' \
                             'PCI_ID %s\n' \
                             '%s_STATUS %s\n' \
                             'HEADER_LOG 0 1 2 3\n\n' % (
                             each_bdf, aer_type, aer_subtype)

        with open('/tmp/aer_file', 'w') as f:
            f.write(file_template)

        out = False
        try:
            if not subprocess.check_call('/usr/bin/aer-inject /tmp/aer_file',
                                         shell=True):
                out = True
        except:
            pass

        return out


    out = aer_inject()
    print "AER injection %s" % 'OK' if out else 'FAIL'
    sys.exit(0)

PCI_DATA = {
    'PCI_COMMAND': {
        'queries':
            [
                {
                    'query': 'I/O',
                    'regex': '[-|+]'
                },
                {
                    'query': 'Mem',
                    'regex': '[-|+]'
                },
                {
                    'query': 'BusMaster',
                    'regex': '[-|+]'
                }
            ],

        'strip': 'Control:'
    },

    'PCI_BUS': {
        'queries':
            [
                {
                    'query': 'primary',
                    'regex': '=\w+'
                },

                {
                    'query': 'secondary',
                    'regex': '=\w+'
                },

                {
                    'query': 'subordinate',
                    'regex': '=\w+'
                },

            ],

        'strip': 'Bus:'

    },

    'PCI_CAP': {
        'queries':
            [
                {
                    'query': 'Enable',
                    'regex': '[-|+]'
                },
                {
                    'query': 'Maskable',
                    'regex': '[-|+]'
                },
                {
                    'query': '64bit',
                    'regex': '[-|+]'
                }
            ],

        'strip': 'MSI:'
    },

    'PCI_EXP_DEVCAP': {
        'queries':
            [
                {
                    'query': 'AttnBtn',
                    'regex': '[-|+]'
                },
                {
                    'query': 'AttnInd',
                    'regex': '[-|+]'
                },
                {
                    'query': 'PwrInd',
                    'regex': '[-|+]'
                }
            ],

        'strip': 'DevCap:'
    },

    'PCI_EXP_DEVSTA': {
        'queries':
            [
                {
                    'query': 'CorrErr',
                    'regex': '[-|+]'
                },
                {
                    'query': 'UncorrErr',
                    'regex': '[-|+]'
                },
                {
                    'query': 'FatalErr',
                    'regex': '[-|+]'
                },
                {
                    'query': 'UnsuppReq',
                    'regex': '[-|+]'
                }
            ],

        'strip': 'DevSta:'
    },

    'PCI_EXP_LNKCAP': {
        'queries':
            [
                {
                    'query': 'Port',
                    'regex': '#\w+'
                },
                {
                    'query': 'Speed',
                    'regex': '\d*.\d*GT/s'
                },
                {
                    'query': 'Width',
                    'regex': 'x\d+'
                },
                {
                    'query': 'ASPM',
                    'regex': '\w+'
                },
                {
                    'query': 'L0',
                    'regex': '<\d+us'
                },
                {
                    'query': 'L1',
                    'regex': '<\d+us'
                }

            ],

        'strip': 'LnkCap:'
    },

    'PCI_EXP_LNKSTA': {
        'queries':
            [
                {
                    'query': 'Speed',
                    'regex': '\d*.\d*GT/s'
                },
                {
                    'query': 'Width',
                    'regex': 'x\d+'
                },
                {
                    'query': 'TrErr',
                    'regex': '[-|+]'
                },
                {
                    'query': 'Train',
                    'regex': '[-|+]'
                },
                {
                    'query': 'SlotClk',
                    'regex': '[-|+]'
                },
                {
                    'query': 'DLActive',
                    'regex': '[-|+]'
                }

            ],

        'strip': 'LnkSta:'
    },

    'PCI_EXP_SLTCAP': {
        'queries':
            [
                {
                    'query': 'AttnBtn',
                    'regex': '[-|+]'
                },
                {
                    'query': 'PwrCtrl',
                    'regex': '[-|+]'
                },
                {
                    'query': 'MRL',
                    'regex': '[-|+]'
                },
                {
                    'query': 'AttnInd',
                    'regex': '[-|+]'
                },
                {
                    'query': 'PwrInd',
                    'regex': '[-|+]'
                },
                {
                    'query': 'HotPlug',
                    'regex': '[-|+]'
                },
                {
                    'query': 'Surprise',
                    'regex': '[-|+]'
                }

            ],

        'strip': 'SltCap:'
    },

    'PCI_EXP_SLTSTA': {
        'queries':
            [
                {
                    'query': 'AttnBtn',
                    'regex': '[-|+]'
                },
                {
                    'query': 'PowerFlt',
                    'regex': '[-|+]'
                },
                {
                    'query': 'MRL',
                    'regex': '[-|+]'
                },
                {
                    'query': 'CmdCplt',
                    'regex': '[-|+]'
                },
                {
                    'query': 'PresDet',
                    'regex': '[-|+]'
                },
                {
                    'query': 'Interlock',
                    'regex': '[-|+]'
                }

            ],

        'strip': 'SltSta:'
    },

    'PCI_EXP_UNCOR_AERMSK': {
        'queries':
            [
                {
                    'query': 'DLP',
                    'regex': '[-|+]'
                },
                {
                    'query': 'SDES',
                    'regex': '[-|+]'
                },
                {
                    'query': 'TLP',
                    'regex': '[-|+]'
                },
                {
                    'query': 'FCP',
                    'regex': '[-|+]'
                },
                {
                    'query': 'CmpltTO',
                    'regex': '[-|+]'
                },
                {
                    'query': 'CmpltAbrt',
                    'regex': '[-|+]'
                },
                {
                    'query': 'UnxCmplt',
                    'regex': '[-|+]'
                },
                {
                    'query': 'RxOF',
                    'regex': '[-|+]'
                },
                {
                    'query': 'MalfTLP',
                    'regex': '[-|+]'
                },
                {
                    'query': 'UnsupReq',
                    'regex': '[-|+]'
                },

            ],

        'strip': 'UEMsk:'
    },

    'PCI_EXP_COR_AERMSK': {
        'queries':
            [
                {
                    'query': 'RxErr',
                    'regex': '[-|+]'
                },
                {
                    'query': 'BadTLP',
                    'regex': '[-|+]'
                },
                {
                    'query': 'BadDLLP',
                    'regex': '[-|+]'
                },
                {
                    'query': 'Rollover',
                    'regex': '[-|+]'
                },
                {
                    'query': 'Timeout',
                    'regex': '[-|+]'
                },
                {
                    'query': 'NotFatalErr',
                    'regex': '[-|+]'
                }
            ],

        'strip': 'CEMsk:'
    },
}

PCI_DATA['PCI_EXP_COR_AERSTA'] = {
    'queries': PCI_DATA['PCI_EXP_COR_AERMSK']['queries'],
    'strip': 'CESta:'
}

PCI_DATA['PCI_EXP_UNCOR_AERSTA'] = {
    'queries': PCI_DATA['PCI_EXP_UNCOR_AERMSK']['queries'],
    'strip': 'UESta:'
}

LIST_QUERIES = PCI_DATA.keys() + ['PCI_MEMORY', 'PCI_NUKE_STATUS',
                                  'PCI_MEMORY_BRIDGE']


def print_query_list():
    print "List of queries"
    print '\n'.join(LIST_QUERIES)


def get_memory(venid, devid, what):
    bdf_list = []
    if options.bdf:
        bdf_list = [options.bdf]
    else:
        if venid and devid:
            lspci = subprocess.check_output('lspci -d %s:%s' % (venid, devid),
                                            shell=True)
            bdf_list = [i.split(' ')[0].strip() for i in lspci.split('\n') if i]
    data = defaultdict()
    for each_bdf in bdf_list:
        out = subprocess.check_output('lspci -s %s -vv' % each_bdf, shell=True)
        pref = '-'
        npref = '-'
        size_pref = '-'
        region_pref = '-'
        size_npref = '-'
        region_npref = '-'
        region = '-'
        size = '-'
        if what == 'PCI_MEMORY_BRIDGE':

            npref = re.search('Memory behind bridge:\s*(\w+)-(\w+).*', out)
            if npref:
                npref = npref.groups()[0] + '-' + npref.groups()[1]
            pref = re.search(
                'Prefetchable memory behind bridge:\s*(\w+)-(\w+).*', out)
            if pref:
                pref = pref.groups()[0] + '-' + pref.groups()[1]
            size = '-'
            region = '-'

            # Instead of handing multiple error codes , we can reuse
            if not npref or not pref:
                print "Something wrong with the bridge or probably a device"
                print "Use -q PCI_MEMORY -d <>"
                return None

        else:
            # Region 2: Memory at b6000000 (32-bit, non-prefetchable) [size=4M]
            memory_found = re.finditer(
                'Region\s*(\d+):\s*Memory at (\w+) \(.*,\s*(.*)\)\s*\[size=(\w+)\]',
                out)
            if memory_found:
                for memory in memory_found:
                    mem_type = memory.groups()[2]
                    if mem_type == 'non-prefetchable':
                        npref = memory.groups()[1]
                        region_npref = memory.groups()[0]
                        size_npref = memory.groups()[3]
                    else:
                        pref = memory.groups()[1]
                        region_pref = memory.groups()[0]
                        size_pref = memory.groups()[3]
                size = '%s/%s' % (size_pref, size_npref)
                region = '%s/%s' % (region_pref, region_npref)
            else:
                print "A bridge mostly"
                print "Use -q PCI_MEMORY_BRIDGE and try again"
                return None

        data[each_bdf] = {
            'pref': pref,
            'npref': npref,
            'size_pref': size_pref,
            'size_npref': size_npref,
            'region_pref': region_pref,
            'region_npref': region_npref,
            'size': '%s/%s' % (size_npref, size_pref),
            'region': '%s/%s' % (region_npref, region_pref)
        }

    return data


def peek_memory(venid, devid, what, offset='0', instance=None):
    if what == 'PCI_MEMORY_BRIDGE':
        print "v1.0 only supports PCI_MEMORY peek"
        sys.exit(1)

    if int(offset, 16) % 4:
        print "Offset should be 4 byte aligned"
        sys.exit(1)

    data = get_memory(venid, devid, what)
    for idx, (each_bdf, val) in enumerate(data.iteritems()):
        # Read the npref first
        printf = 'printf "%x"'
        # Just making sure we have some address
        if len(val['pref']) >= 9:
            added = subprocess.check_output('%s $((0x%s + 0x%s))' %
                                            (printf, val['pref'],
                                             offset.replace('0x', '').replace(
                                                 '0X', '')),
                                            shell=True)
            npref_w_off = added.strip().replace('\n', '')
            value = subprocess.check_output(
                'pcimemread %s 4 | grep -v PCI | cut -d ":" -f2' % npref_w_off,
                shell=True).strip().replace('\n', '').rstrip('.').replace(' ',
                                                                          '')
            print "pref:%s:%s:%s:%s" % (
            idx, npref_w_off, value, 'FAIL' if 'ffffffff' in value else 'OK')

        if len(val['npref']) == 8:
            added = subprocess.check_output('%s $((0x%s + 0x%s))' %
                                            (printf, val['npref'],
                                             offset.replace('0x', '').replace(
                                                 '0X', '')),
                                            shell=True)
            pref_w_off = added.strip().replace('\n', '')
            value = subprocess.check_output(
                'pcimemread %s 4 | grep -v PCI | cut -d ":" -f2' % pref_w_off,
                shell=True).strip().replace('\n', '').rstrip('.').replace(' ',
                                                                          '')
            print "npref:%s:%s:%s:%s" % (
            idx, pref_w_off, value, 'FAIL' if 'ffffffff' in value else 'OK')


if options.query_list:
    print_query_list()
    sys.exit(0)

venid = None
devid = None
if options.vendev:
    vendev = options.vendev.split(':')
    if len(vendev) == 2:
        venid = vendev[0]
        devid = vendev[1]
    else:
        print "Format for devid and venid is -d 1137:00ca <venid>:<devid>"
        sys.exit(1)

elif options.bdf:
    bdf_check = re.search('(\w\w):(\w\w).(\w)', options.bdf)
    if not bdf_check:
        print "Format for bdf is bb:dd.f"
        sys.exit(1)

else:
    print "Use -d 1137:00ca OR use -s bb:dd.f as you would normally use in lspci"
    sys.exit(1)

if options.query and options.query not in LIST_QUERIES:
    print "Query not in list"
    print_query_list()
    sys.exit(1)

SPACE = ' '

if options.query == 'PCI_NUKE_STATUS':
    if not options.raw:
        print "-----------------------------------------------------------------"
        print 'bdf' + SPACE * 7 + 'nukes' + SPACE * 2 + 'linkdown' + SPACE * 2 + 'last-nuke' + SPACE * 13 + 'last-link' + \
              SPACE * 13 + 'claimed-by'
        print "-----------------------------------------------------------------"
    if options.bdf:
        bdf_list = [options.bdf]
    else:
        lspci = subprocess.check_output('lspci -d %s:%s' % (venid, devid),
                                        shell=True)
        bdf_list = [i.split(' ')[0].strip() for i in lspci.split('\n') if i]
    for each_bdf in bdf_list:
        try:
            claims = [i.strip() for i in subprocess.check_output(
                'dmesg | grep "%s" | grep claimed' % each_bdf,
                shell=True).split('\n') if i]
        except subprocess.CalledProcessError:
            claims = []
        try:
            nukes = [i.strip() for i in subprocess.check_output(
                'dmesg | grep "%s" | grep nuke | grep MMIO' % each_bdf,
                shell=True).split('\n') if i]
        except subprocess.CalledProcessError:
            nukes = []
        try:
            linkdowns = [i.strip() for i in subprocess.check_output(
                'dmesg | grep "%s" | grep "Link Down"' % each_bdf,
                shell=True).split('\n') if i]
        except subprocess.CalledProcessError:
            linkdowns = []

        claim = len(claims)
        nuke = len(nukes)
        link = len(linkdowns)
        by = '-'
        last_nuke = '-'
        last_ld = '-'
        claimed_by = '-'
        if claims:
            claimed_by = re.search('.*claimed by\s*(\w+).*', claims[-1])
            if claimed_by:
                by = claimed_by.groups()[0]
        if nukes:
            last_nuke = get_timestamp(nukes[-1])
        if linkdowns:
            last_ld = get_timestamp(linkdowns[-1])

        if options.raw:
            print str(each_bdf) + ": nukes: " + str(nuke)
            print str(each_bdf) + ": linkdown: " + str(link)
            print str(each_bdf) + ": last-nuke: " + str(last_nuke)
            print str(each_bdf) + ": last-link: " + str(last_ld)
            print str(each_bdf) + ": claimed-by: " + str(by)
        else:
            print str(each_bdf) + (SPACE * 3) + str(nuke) + (SPACE * 7) + str(
                link) + (SPACE * 10) + str(last_nuke) + \
                  (SPACE * (14 - len(last_nuke))) + str(last_ld) + (
                  SPACE * (14 - len(last_ld))) + str(by)

    if not options.raw:
        print "-----------------------------------------------------------------"
    sys.exit(0)

if not options.query:
    print "Need query use -l to get list of queries and use -q to name a query"
    sys.exit(1)

final = defaultdict(dict)
what = options.query

if 'PCI_MEMORY' in what:

    if options.peek:
        peek_memory(venid, devid, what, offset=options.peek)

    else:
        if not options.raw:
            print "-----------------------------------------------------------------"
            if what == 'PCI_MEMORY_BRIDGE':
                print 'bdf' + SPACE * 7 + 'non-pref' + SPACE * 17 + 'pref' + SPACE * 33 + 'size' + SPACE + 'region'
                print "-----------------------------------------------------------------"
            else:
                print 'bdf' + SPACE * 7 + 'non-pref' + SPACE * 16 + 'pref' + SPACE * 8 + 'size' + SPACE + 'region'
                print "-----------------------------------------------------------------"

        data = get_memory(venid, devid, what)
        if not data:
            print "Failure"
            sys.exit(1)

        for each_bdf, val in data.items():
            if what == 'PCI_MEMORY_BRIDGE':

                if options.raw:
                    print each_bdf + " :npref: " + val['npref']
                    print each_bdf + " :pref: " + val['pref']
                    print each_bdf + " :size: " + val['size']
                    print each_bdf + " :region: " + val['region']
                else:
                    print each_bdf + SPACE + val['npref'] + SPACE + val[
                        'pref'] + SPACE + val['size'] + SPACE + val['region']
            else:
                if options.raw:
                    print each_bdf + " :npref: " + val['npref']
                    print each_bdf + " :pref: " + val['pref']
                    print each_bdf + " :size: " + val['size']
                    print each_bdf + " :region: " + val['region']
                else:
                    print each_bdf + SPACE + val['npref'] + SPACE * 11 + val[
                        'pref'] + SPACE * 2 + val['size'] + SPACE + \
                          val['region']

        if not options.raw:
            print "-----------------------------------------------------------------"


else:
    if not options.raw:
        print "-----------------------------------------------------------------"
        print 'bdf' + SPACE * 7 + SPACE.join(
            [i['query'] for i in PCI_DATA[what]['queries']])
        print "-----------------------------------------------------------------"
    if options.bdf:
        bdf_list = [options.bdf]
    else:
        lspci = subprocess.check_output('lspci -d %s:%s' % (venid, devid),
                                        shell=True)
        bdf_list = [i.split(' ')[0].strip() for i in lspci.split('\n') if i]
    for each_bdf in bdf_list:
        out = subprocess.check_output('lspci -s %s -vv' % each_bdf, shell=True)
        if options.raw:
            out_str = ''
        else:
            out_str = each_bdf + SPACE + SPACE * len('bdf')
        filtered = re.search('%s\s*(.*)\n' % PCI_DATA[what]['strip'], out)
        if not filtered:
            continue

        for each in PCI_DATA[what]['queries']:
            a = re.search('%s\s*(%s)' % (each['query'], each['regex']),
                          filtered.groups()[0], re.I)
            if not a:
                continue

            if options.raw:
                out_str += "%s : %s : %s\n" % (
                each_bdf, each['query'], a.groups()[0].replace('=',''))
            else:
                if a:
                    out_str += a.groups()[0].replace('=','') + SPACE * (
                    len(each['query']) - len(a.groups()[0].replace('=','')) + 3)
                else:
                    out_str += SPACE * (len(each['query']) - len(a.groups()[0]) + 3)

        print out_str

    if not options.raw:
        print "-----------------------------------------------------------------"
