#!/usr/bin/perl

use strict;
use Getopt::Long;

my $idsfile = "/etc/recoverycd/pci.ids";
my $tablefile = "/etc/recoverycd/pcitable";

sub usage
{
	print STDOUT <<END;
Usage:
	rcdProbePCI [-h | --help | --show-detail]
END

}
#############################################################################
# Param:
#		<Vendor ID>:<Product ID>:<Sub Vendor ID>:<Sub Product ID>
# Return:
#		name of device
#############################################################################
sub getNameByID
{
	my ($vendorID, $productID, $subVendorID, $subProductID)  = split(/:/, $_[0]);
	
	my $vendorName;
	my $productName;
	my $subSysName;
	
	open FILE, $idsfile;
	my @lines = <FILE>;
	close FILE;

	my $deep = 0;
	foreach ( @lines ) {
		chomp;
		# skip comment, empty, and "C " lines
		if (/[\t]*#/ && /^\s*$/ && /^C /) {
			next;
		}
		# exit if out of range
		my $tabnum = $deep - 1;
		if ($deep > 0 && /^\t{0,$tabnum}\w/) {
			last;
		}
		# vendor line
		if ($deep == 0 && s/^$vendorID[\t ]*//) {
			$vendorName = $_;
			$deep = 1;
			next;
		}
		# product line
		if ($deep == 1 && s/^\t$productID[\t ]*//) {
			$productName =  $_;
			$deep = 2;
			next;
		}
		# sub system line
		if ($deep == 2 && s/^\t\t$subVendorID $subProductID[\t ]*//) {
			$subSysName = $_;
			last;
		}
	}
	
	if(!defined $vendorName) {
		return "PCI device $vendorID:$productID";
	} elsif(!defined $productName) {
		return "PCI device $vendorID:$productID ($vendorName)";
	} else {
		return "$vendorName $productName $subSysName";
	}
}

#############################################################################
# Param:
#	<Vendor ID> <Product ID> <Class>
# Return:
#	1. if the device ID match the table and the device is configed in
#          raid mode, return 1.
#       2. else, return 0
#############################################################################
my @adpahci_maps = (
	{ VID => "8086", DID => "2652", },
	{ VID => "8086", DID => "2682", },
	{ VID => "8086", DID => "27c3", }
);

sub CheckAdpahciModule
{
	my $vendorID = shift;
	my $productID = shift;
	my $class = shift;
	my $match = 0;
	my $ret = 0;

	foreach my $device ( @adpahci_maps ) {
		if ($device->{'VID'} eq "$vendorID" && $device->{'DID'} eq "$productID") {
			$match = 1;
			last;
		}
	}

	if ($match == 1) {
		my $subclass = substr($class, 2, 2);
		$ret = 1 if ($subclass eq "04");
	}

	return $ret;
}

my @specahci_maps = (
	{ VID => "8086", DID => "2652", },
	{ VID => "8086", DID => "2653", },
);

sub CheckSpecahciModule
{
	my $vendorID = shift;
	my $productID = shift;
	my $class = shift;
	my $match = 0;
	my $ret = 0;

	foreach my $device ( @specahci_maps ) {
		if ($device->{'VID'} eq "$vendorID" && $device->{'DID'} eq "$productID") {
			$match = 1;
			last;
		}
	}

	if ($match == 1) {
		my $subclass = substr($class, 2, 2);
		$ret = 1 if ($subclass eq "01");
	}

	return $ret;
}

my @aar81xx_module_maps = (
	{ VID => "11ab", DID => "5041", },
	{ VID => "11ab", DID => "5081", },
	{ VID => "11ab", DID => "6041", },
	{ VID => "11ab", DID => "6081", },
	{ VID => "9005", DID => "0241", },
	{ VID => "9005", DID => "0243", },
);

sub check_aar81xx_module
{
	my $vendorID = shift;
	my $productID = shift;
	my $class = shift;
	my $match = 0;
	my $ret = 0;

	foreach my $device ( @aar81xx_module_maps ) {
		if ($device->{'VID'} eq "$vendorID" && $device->{'DID'} eq "$productID") {
			$match = 1;
			last;
		}
	}

	if ($match == 1) {
		my $subclass = substr($class, 2, 2);
		$ret = 1 if ($subclass eq "04");
	}

	return $ret;
}

#############################################################################
# Param:
#		<Vendor ID>:<Product ID>:<Sub Vendor ID>:<Sub Product ID>
# Return:
#		1. if the driver is a module, return the name of module
#		2. if the driver is builtin, return ""
#		3. if the driver not exists, return undef
#############################################################################
sub getModuleByID
{
	my ($vendorID, $productID, $subVendorID, $subProductID)  = split(/:/, $_[0]);
	
	open FILE, $tablefile;
	my @lines = <FILE>;
	close FILE;
	
	my $module;
	
	my $line;
	foreach $line(@lines) {
		chomp $line;
		# skip empty line and comment line
		if(length $line == 0 || substr($line, 0, 1) eq "#" ) {
			next;
		}
		# split
		my @words = split(/[\t ]+/, $line);
		# without sub system
		if(@words == 3) {
			if($words[0] eq "0x$vendorID" && $words[1] eq "0x$productID") {
				$module =  $words[2];
			}
		} elsif(@words == 5) {
			if($words[0] eq "0x$vendorID" && $words[1] eq "0x$productID"
				&& $words[2] eq "0x$subVendorID" && $words[3] eq "0x$subProductID" ) {
				$module =  $words[4];
			}
		}
	}

	return substr($module, 1, length($module)-2);
}

#############################################################################
# Return:
#		refrence of pci devices array
#############################################################################
sub getPCIDevices
{
	my @devices;
	
	# get all buses
	opendir DIR, "/proc/bus/pci";
	my @buses = grep(/^[0-9a-f]{2}$/, readdir DIR);
	closedir DIR;
	
	# scan each bus
	my $bus;
	foreach $bus (@buses) {
		opendir DIR, "/proc/bus/pci/$bus";
		my @entries = grep(/^[0-9a-f]{2}\.[0-9a-f]/, readdir DIR);
		closedir(DIR);
		
		my $entry;
		foreach $entry (@entries) {
			my ($dev, $func) = split(/\./, $entry);
			
			open FILE, "/proc/bus/pci/$bus/$entry";
			binmode(FILE);
			read(FILE, my $data, 64);
			close(FILE);
			
			my %device;
			$device{'Bus'} = $bus;
			$device{'DeviceID'} = $dev;
			$device{'FunctionID'} = $func;
			
			$device{'VendorID'} = sprintf "%02x%02x", vec($data, 1, 8), vec($data, 0, 8);
			$device{'ProductID'} = sprintf "%02x%02x", vec($data, 3, 8), vec($data, 2, 8);
			$device{'Revision'} = sprintf "%02x", vec($data, 8, 8);
			$device{'ClassCode'} = sprintf "%02x%02x%02x", vec($data, 11, 8), vec($data, 10, 8), vec($data, 9, 8);
			$device{'SubVendorID'} = sprintf "%02x%02x", vec($data, 45, 8), vec($data, 44, 8);
			$device{'SubProductID'} = sprintf "%02x%02x", vec($data, 47, 8), vec($data, 46, 8);
			
			$devices[@devices] = \%device;
		}
	}
	
	return \@devices;
}
#############################################################################
# Param:
#		the module name
# Return:
#		0 if load module successfully, or the module is also loaded
#		non 0 if load module failed
#############################################################################
sub loadModule
{
	my $module = shift;
	my $ret;
	
	system("modprobe $module &> /dev/null");
	if ($? != 0) {
		$ret = 1;
	}
	else {
		$ret = system("lsmod | grep \"^$module \" &> /dev/null") >> 8;
	}
	
	return $ret;
}
#############################################################################
# Main program
#############################################################################
my $devices = getPCIDevices();
my $device;

my @installed_devices;
my @failed_devices;
my @unknown_devices;

my $help_flag;
my $show_detail_flag;

GetOptions('h|help' => \$help_flag, 'show-detail' => \$show_detail_flag);

# show help message only
if($help_flag) {
	usage();
	exit 0;
}
# probe all pci device
foreach $device(@{$devices}) {
	my $baseClass = substr($device->{'ClassCode'}, 0, 2);
	my $class = substr($device->{'ClassCode'}, 0, 4);
	# 01: Storage, 02: Network, 0c04: Fibre Channel
	if($baseClass ne "01" && $baseClass ne "02" && $class ne "0c04") {
		next;
	}
	# load module
	
	my $id = $device->{'VendorID'}.":".$device->{'ProductID'}.":".$device->{'SubVendorID'}.":".$device->{'SubProductID'};
	my $module = getModuleByID($id);
	my $name = getNameByID($id)." (rev ".$device->{'Revision'}.")";

	# check adpahci module
	if (CheckAdpahciModule($device->{'VendorID'}, $device->{'ProductID'}, $device->{'ClassCode'}) == 1) {
		$module = "adpahci";
	}
	elsif (CheckSpecahciModule($device->{'VendorID'}, $device->{'ProductID'}, $device->{'ClassCode'}) == 1) {
		$module = "ata_piix";
	}
	elsif (check_aar81xx_module($device->{'VendorID'}, $device->{'ProductID'}, $device->{'ClassCode'}) == 1) {
		$module = "aar81xx";
	}

	#unknown device
	if(!defined $module) {
		$unknown_devices[@unknown_devices] = $name."($id)";
	}
	# builtin device
	elsif($module eq "") {
		$installed_devices[@installed_devices] = $name;
	}
	# module device
	else {
		my $ret = loadModule($module);
		if($ret == 0) {
			$installed_devices[@installed_devices] = $name;
		} else {
			$failed_devices[@failed_devices] = $name."($id)";
		}
	}
}
# show result
if(!$show_detail_flag) {
	foreach $device (@installed_devices) {
		print STDOUT "$device\n";
	}
} else {
	# show installed devices
	print STDOUT "Installed devices:\n";
	if(@installed_devices){
		foreach $device (@installed_devices) {
			print STDOUT "$device\n";
		}
	}else{
		print STDOUT "None\n";
	}
	print STDOUT "\n";
	# show failed devices
	print STDOUT "Failed devices:\n";
	if(@failed_devices){
		foreach $device (@failed_devices) {
			print STDOUT "$device\n";
		}
	}else{
		print STDOUT "None\n";
	}
	print STDOUT "\n";
	# show unknown devices
	
	print STDOUT "Unknown devices:\n";
	if(@unknown_devices){
		foreach $device (@unknown_devices) {
			print STDOUT "$device\n";
		}
	}else{
		print "None\n";
	}
}
