###########################################################################
# ScanPCI.py -                                                            #
# ------------------------------                                          #
# copyright : (C) 2005 by Simon Edwards                                   #
# email     : simon@simonzone.com                                         #
#                                                                         #
###########################################################################
#                                                                         #
#   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.                                   #
#                                                                         #
###########################################################################
"""Provides information about the devices attached to the PCI bus.
"""
import struct
import csv
import os.path
import sys

###########################################################################
class PCIDevice(object):
    def __init__(self,line=None):
        self.vendor = None      # PCI vendor id
        self.device = None

        self.subvendor = None   # 0xffff if not probe_type'd or no subid
        self.subdevice = None   # 0xffff if not probe_type'd or no subid
        self.pci_class = None   # 'None' if not probe_type'd

        self.pci_bus = None     # pci bus id 8 bits wide
        self.pci_device = None  # pci device id 5 bits wide
        self.pci_function = None# pci function id 3 bits wide

        self.module = None
        self.text = None
        self.already_found = False

        if line is not None:
            self.loadFromString(line)

    def isGfxCard(self):
        if self.module is not None and \
                (self.module.startswith("Card:") or self.module.startswith("Server:XFree86(")):
            return True

        return (self.pci_class & PCIBus.PCI_BASE_CLASS_MASK)==PCIBus.PCI_BASE_CLASS_DISPLAY

    def getModule(self):
        if self.module is not None:
            if self.module.startswith("Server:XFree86("):
                return self.module[15:-1]
            elif self.module.startswith("Card:"):
                return self.module[5:]
        return self.module

    def isModuleXorgDriver(self):
        return self.module is not None and \
            (self.module.startswith("Server:XFree86(") or self.module.startswith("Card:"))

    def __str__(self):
        s = "PCI:%i:%i:%i, " % (self.pci_bus,self.pci_device,self.pci_function)
        s += "Vendor:%x, Device:%x," % (self.vendor,self.device)
        if self.subvendor is not None:
            s += " Subvendor:%x," % self.subvendor
        if self.subdevice is not None:
            s += " Subdevice:%x," % self.subdevice
        if self.pci_class is not None:
            s += " Class:%x," % self.pci_class
        if self.module is not None:
            s += " Module:%s," % self.module
        if self.text is not None:
            s += " Text:%s" % self.text
        return s

    def loadFromString(self,line):
        parts = line.split(",")
        for i in range(len(parts)):
            bit = parts[i].strip()
            if bit.startswith("PCI:"):
                pci_code = bit[4:].split(":")
                self.pci_bus = int(pci_code[0])
                self.pci_device = int(pci_code[1])
                self.pci_function = int(pci_code[2])
            elif bit.startswith("Vendor:"):
                self.vendor = int(bit[7:],16)
            elif bit.startswith("Device:"):
                self.device = int(bit[7:],16)
            elif bit.startswith("Subvendor:"):
                self.subvendor = int(bit[10:],16)
            elif bit.startswith("Subdevice:"):
                self.subdevice = int(bit[10:],16)
            elif bit.startswith("Class:"):
                self.pci_class = int(bit[6:],16)
            elif bit.startswith("Module:"):
                self.module = bit[7:]
            elif bit.startswith("Text:"):
                self.text = " ".join(parts[i:]).strip()[5:]
                break

############################################################################
class PCIBus(object):
    PCI_CLASS_SERIAL_USB = 0x0c03
    PCI_CLASS_SERIAL_FIREWIRE = 0x0c00
    PCI_BASE_CLASS_MASK = 0xff00
    PCI_BASE_CLASS_DISPLAY = 0x0300

    def __init__(self, data_file_dir="."):
        self.devices = []
        self.data_file_dir = data_file_dir

    def detect(self,device_data="/proc/bus/pci/devices"):
        # Shamelessly translated from ldetect's pci.c.
        fhandle = open(device_data)
        for line in fhandle.readlines():
            #print "L:",line
            entry = PCIDevice()
            self.devices.append(entry)
            parts = line.split()

            devbusfn = int(parts[0],16)
            idbits = int(parts[1],16)
            entry.vendor = idbits >> 16
            entry.device = idbits & 0xffff
            entry.pci_bus = devbusfn >> 8
            entry.pci_device = (devbusfn & 0xff) >> 3
            entry.pci_function = (devbusfn & 0xff) & 0x07

            try:
                infohandle = open("/proc/bus/pci/%02x/%02x.%d" % (
                                entry.pci_bus, entry.pci_device, entry.pci_function),"r")
                # these files are 256 bytes but we only need first 48 bytes
                buf = infohandle.read(48)   
                (class_prog, entry.pci_class, entry.subvendor, entry.subdevice) = \
                        struct.unpack("<xxxxxxxxxBHxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxHH",buf)
                #print "STRUCT: ",struct.unpack("@xxxxxxxxxBHxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxHH",buf)
                if (entry.subvendor==0 and entry.subdevice==0) or \
                        (entry.subvendor==entry.vendor and entry.subdevice==entry.device):
                    entry.subvendor = 0xffff
                    entry.subdevice = 0xffff
                if entry.pci_class == PCIBus.PCI_CLASS_SERIAL_USB:
                    # taken from kudzu's pci.c
                    if class_prog == 0:
                        entry.module = "usb-uhci"
                    elif class_prog == 0x10:
                        entry.module = "usb-ohci"
                    elif class_prog == 0x20:
                        entry.module = "ehci-hcd"
                if entry.pci_class == PCIBus.PCI_CLASS_SERIAL_FIREWIRE:
                    # taken from kudzu's pci.c
                    if class_prog == 0x10:
                        entry.module = "ohci1394"
                infohandle.close()
            except IOError:
                pass
        fhandle.close()

        #if False or os.path.exists("/usr/share/ldetect-lst/pcitable"):
        #self._resolveDevicesWithLdetect()
        #else:
        self._resolveDevicesWithHwdata()
        #self._resolveDevicesWithDiscover()

    def _resolveDevicesWithLdetect(self):
        # Scan the PCI database.
        #fhandle = open(os.path.join(self.data_file_dir,"pcitable"),"r")
	fhandle = open(os.path.join("/opt/trinity/share/apps/guidance/","pcitable"),"r")

        # This class is just for skipping comment lines in the database file.
        # This whole class is just an iterator wrapper that we put around our file iterator.
        class commentskipperiterator(object):
            def __init__(self,fhandle):
                self.fhandle = iter(fhandle)
            def __iter__(self):
                return self
            def next(self):
                line = self.fhandle.next()
                while line[0]=="#":
                    line = self.fhandle.next()
                return line

        unknowndevices = self.devices[:]

        # Process each row of the DB.
        for row in csv.reader(commentskipperiterator(fhandle),delimiter='\t'):
            if len(row)==4:
                (vendor,device,module,text) = row
            elif len(row)==6:
                (vendor, device, subvendor, subdevice, module, text) = row
                subvendor = int(subvendor[2:],16)
                subdevice = int(subdevice[2:],16)
            else:
                continue
            vendor = int(vendor[2:],16) # parse hex numbers of the form 0x1abc
            device = int(device[2:],16)

            i = 0
            while i<len(unknowndevices):
                pcidevice = unknowndevices[i]
                if pcidevice.vendor==vendor and pcidevice.device==device \
                        and (len(row)==4 \
                        or (pcidevice.subvendor==subvendor and pcidevice.subdevice==subdevice)):
                    if module!="unknown":
                        pcidevice.module = module
                    pcidevice.text = text
                    if len(row)==6: # Close match, also matched on subdevice/subvendor ids.
                        del unknowndevices[i]
                    else:
                        i += 1
                else:
                    i += 1

        fhandle.close()

    def _resolveDevicesWithDiscover(self):

        unknown_devices = self.devices[:]
        self._resolveDevicesWithDiscoverFile("/usr/share/discover/pci-26.lst",unknown_devices)
        self._resolveDevicesWithDiscoverFile("/usr/share/discover/pci.lst",unknown_devices)

    def _resolveDevicesWithDiscoverFile(self,filename,unknown_devices):
        # Scan the PCI database.
        fhandle = open(filename,"r")

        # Process each row of the DB.
        for line in fhandle:
            row = line.replace("\t"," ").split(" ")
            if len(row) >= 1 and row[0] != '':
                # Skip manufacturer info lines.
                continue

            vendor = int(row[1][:4],16)
            device = int(row[1][4:],16)
            module = row[3]
            text = ' '.join(row[4:]).strip()

            i = 0
            while i<len(unknown_devices):
                pcidevice = unknown_devices[i]
                if pcidevice.vendor==vendor and pcidevice.device==device:
                    pcidevice.module = module
                    pcidevice.text = text
                    del unknown_devices[i]
                else:
                    i += 1

        fhandle.close()

    def _resolveDevicesWithHwdata(self):
        # Scan the PCI database.
        fhandle = open("/usr/share/hwdata/pci.ids","r")

        # This class is just for skipping comment lines in the database file.
        # This whole class is just an iterator wrapper that we put around our file iterator.
        class commentskipperiterator(object):
            def __init__(self,fhandle):
                self.fhandle = iter(fhandle)
            def __iter__(self):
                return self
            def next(self):
                line = self.fhandle.next()
                while line[0]=="#":
                    line = self.fhandle.next()
                return line

        unknowndevices = self.devices[:]

        # Process each row of the DB.
        for row in fhandle:
            stripped_row = row.strip()

            if stripped_row=='' or stripped_row[0]=='#':
                continue # Comment or blank line, skip it.

            if stripped_row[0]=='C':
                # Reached the device class data, stop.
                break

            if row[0]!='\t':
                # Vendor line
                vendor_parts = stripped_row.split(' ')
                vendor = int(vendor_parts[0],16)
                continue

            if row[1]!='\t':
                # Device line
                device_parts = stripped_row.split(' ')
                device = int(device_parts[0],16)
                subvendor = None
                subdevice = None
            else:
                # Subvendor line
                subvendor_parts = stripped_row.split(' ')
                subvendor = int(subvendor_parts[0],16)
                subdevice = int(subvendor_parts[1],16)

            i = 0
            while i<len(unknowndevices):
                pcidevice = unknowndevices[i]
                if pcidevice.vendor==vendor and pcidevice.device==device \
                        and (subvendor is None \
                        or (pcidevice.subvendor==subvendor and pcidevice.subdevice==subdevice)):
                    #pcidevice.module = module
                    if subvendor is None:
                        pcidevice.text = ' '.join(vendor_parts[1:]) + '|' + ' '.join(device_parts[1:]).strip()
                        i += 1
                    else:
                        pcidevice.text = ' '.join(vendor_parts[1:]) + '|' + ' '.join(device_parts[1:]+subvendor_parts[2:]).strip()
                        del unknowndevices[i] # Perfect match, finished with this device.
                else:
                    i += 1

        fhandle.close()

    def __str__(self):
        return "\n".join([str(x) for x in self.devices])

    def loadFromFile(self,filename):
        fhandle = open(filename,'r')
        for line in fhandle.readlines():
            if line.strip()!="":
                entry = PCIDevice(line=line)
                self.devices.append(entry)
        fhandle.close()

############################################################################
def main():
    bus = PCIBus("ldetect-lst/")
    if len(sys.argv)>1:
        if sys.argv[1]=="--help" or sys.argv[1]=="-h":
            print "Usage:\n    ScanPCI.py <pci device file name>"
            sys.exit(0)
        bus.detect(sys.argv[1])
    else:
        bus.detect()
    print bus

if __name__=='__main__':
    main()