#!/usr/bin/perl -w

use strict;
use Getopt::Long;

my $CONFIG = "/etc/mdadm/mdadm.conf";
my $STATE = "/proc/mdstat";
my $SCSI = "/proc/scsi/scsi";
my $SCSIINFO = "/usr/sbin/scsiinfo";
my $SFDISK = "/sbin/sfdisk";
my $MDADM = "/sbin/mdadm";
my $GRUB = "/sbin/grub";

my @BOOTDRIVES = ("sda", "sdb");
my $TMPPATH = "/root/fixraid";
my $TMPSFDISK = "$TMPPATH.sfdisk.";
my $TMPSCSI = "$TMPPATH.scsi.";
my $TMPGRUB = "$TMPPATH.grub.";

my %mdConfig;
my %mdState;
my %scsi;
my %args;

my ($remove, $add, $fix, $test);

GetOptions('remove=s'	=> \$remove,
	   'add=s'	=> \$add,
	   'fix'	=> \$fix);

SanityCheck();
ReadConfig();
ReadState();
ReadSCSI();

if (defined($remove)) {
	Remove($remove);
} elsif (defined($add)) {
	Add($add);
} elsif (defined($fix)) {
	Fix($fix);
} else {
	ShowFailed();
}

Exit(0);

sub ReadConfig {
	# Build matrix of md's and the sd's that are supposed to be 
	# part of them. Store RAID level as hash value
	my ($md, $raid, @parts);
	undef(%mdConfig);
	open(CONFIG, $CONFIG);
	while (my $line = <CONFIG>) {
		if ($line =~ m/^ARRAY (\S+) level=(raid\d+) .*$/) {
			$md = $1;
			$raid = $2;
			$line = <CONFIG>;
			$line =~ s/^\s+devices=(.*)\n$/$1/;
			@parts = split(/,/, $line);
			@parts = sort(@parts);
			foreach my $sd (@parts) {
				$mdConfig{$md}{$sd} = $raid;
			}
		}
	}
	close(CONFIG);
}

sub ReadState {
	# Build matrix of md's and the sd's that are actually present, 
	# noting their state. (sd could be present but failed.)
	my ($md, @parts);
	undef(%mdState);
	open(STATE, $STATE);
	while (my $line = <STATE>) {
		if ($line =~ m/^(md\d+).*$/) {
			$md = "/dev/$1";
			if ($line =~ m/^.* active raid\d+ (.*)\n$/) {
				$line = $1;
				@parts = split(/\s+/, $line);
				@parts = sort(@parts);
				foreach my $sd (@parts) {
					if ($sd =~ m/\(F\)/) {
						$sd =~ s/\[.*$//;
						$sd = "/dev/$sd";
						$mdState{$md}{$sd} = "F";
					} else {
						$sd =~ s/\[.*$//;
						$sd = "/dev/$sd";
						$mdState{$md}{$sd} = "U";
					}
				}
			} else {
				die("ReadState() confused by:\n$line");
			}
		}
	}
	close(STATE);
}

sub ReadSCSI {
	# Build hash of drive letter ID's and numeric SCSI paths
	# Devices may not be in order in /proc/scsi/scsi
	my (@drives, @devices, $bus, $chan, $id, $lun);
	my $i = "0";

	undef(%scsi);

	open(SCSI, $SCSI);
	local $/ = "Host:";
	while (<SCSI>) {
		s/.*?
		    scsi	(\d+).+?
		    Channel:\s+	(\d+).+?
		    Id:\s+	(\d+).+?
		    Lun:\s+	(\d+).+?
		    Type:\s+	(\S+).*
		/$1 $2 $3 $4/sgx;
		push(@devices, $_) if (defined($5) && $5 eq 'Direct-Access');
	}
	close(SCSI);

	@drives = split(' ', qx($SCSIINFO -l));
	@devices = sort(@devices); 
	foreach my $path (@devices) {
		$scsi{$drives[$i]} = $path;
		$i++;
	}
}

sub ShowFailed {
	# Display md's with missing or failed sd's, and RAID level
	my $fail = 0;
	foreach my $md (keys(%mdConfig)) {
		foreach my $sd (keys(%{$mdConfig{$md}})) {
			if (IsDown($sd) || IsFailed($sd)) {
				print("$md $mdConfig{$md}{$sd} $sd\n");
				$fail = 1;
			}
		}
	}
	Exit(1) if ($fail);
}

sub Failures {
	ReadState();
	my @failures;
	foreach my $md (keys(%mdConfig)) {
		foreach my $sd (keys(%{$mdConfig{$md}})) {
			if (IsDown($sd) || IsFailed($sd)) {
			    	push(@failures, $sd);
			}
		}
	}
	return(@failures);
}
		

sub Remove {
	# Remove the specified sd from any md's it is part of, store SCSI path
	# and partition information, and remove the sd from /proc/scsi/scsi
	my $drive = shift();

	if (!IsPresentSCSI($drive)) {
		Exit(1, "/dev/$drive is not present in system");
	}

	PreRemoveCheck($drive);
	
	foreach my $md (keys(%mdState)) {
		foreach my $sd (keys(%{$mdState{$md}})) {
			if ($sd =~ m/$drive/) {
				if (IsUp($sd)) {
					RAIDSetFaulty($md, $sd);
				}
				RAIDHotRemove($md, $sd);
			}
		}
	}

	SaveDrive($drive);

	ReadState();
	if (IsPresentRAID($drive)) {
		Exit(1, "/dev/$drive RAID removal failed");
	}

	# Don't iterate through this hash, we know the hash key - ARGV STANDARD
	foreach my $sd (keys(%scsi)) {
		if ($sd =~ m/$drive/) {
			system("echo 'scsi remove-single-device $scsi{$sd}' " .
			    "> /proc/scsi/scsi");
			last;
		}
	}

	ReadSCSI();
	if (IsPresentSCSI($drive)) {
		Exit(2, "/dev/$drive SCSI removal failed");
	}
}

sub PreRemoveCheck {
	# Make sure drive to be removed is not part of a RAID0 array,
	# would not be the second missing element in a RAID5 array,
	# and isn't the last element in a RAID1 array.
	# 
	# This function could stand a rewrite.
	#
	my $drive = shift();

	foreach my $md (keys(%mdConfig)) {
		foreach my $sd (keys(%{$mdConfig{$md}})) {
			if ($sd =~ m/$drive/) {
				if ($mdConfig{$md}{$sd} eq 'raid0') {
					Exit(2, "Cannot remove /dev/$drive:\n" .
					    "$sd is device of raid0 array $md");
				} elsif ($mdConfig{$md}{$sd} eq 'raid1') {
					if (keys(%{$mdState{$md}}) < 2 && IsUp($sd)) {
						Exit(2, "Cannot remove /dev/$drive:\n" .
						    "$sd is final device of raid1 array $md");
					}
				} elsif ($mdConfig{$md}{$sd} eq 'raid5') {
					if ((keys(%{$mdConfig{$md}}) - keys(%{$mdState{$md}})) == 1 && IsUp($sd)) {
						Exit(2, "Cannot remove /dev/$drive:\n" .
						    "$sd is active device of raid5 array $md which is already missing one device.");
					}
				} else {
					Exit(2, "/dev/$drive is part of $mdConfig{$md}{$sd} array $md\n" .
					    "Don't know how to operate on $mdConfig{$md}{$sd} arrays");
				}
			}
		}
	}
}

sub Add {
	# Insert the specified sd in /proc/scsi/scsi, repartition, 
	# and insert sd into all appropriate md's
	my $drive = shift();

	if (!IsPresentSCSI($drive)) {
		ReinsertDrive($drive);
	}
	
	foreach my $md (keys(%mdConfig)) {
		foreach my $sd (keys(%{$mdConfig{$md}})) {
			if ($sd =~ m/$drive/) {
				if (IsFailed($sd)) {
					RAIDHotRemove($md, $sd);
					RAIDHotAdd($md, $sd);
				} elsif (IsDown($sd)) {
					RAIDHotAdd($md, $sd);
				}
			}
		}
	}

	ReadState();
	if (!IsPresentRAID($drive)) {
		Exit(2, "/dev/$drive RAID add failed");
	}
	Grub();
}

sub Fix {
	my ($md, $sd);
	foreach $sd (Failures()) {
		($md) = MemberOf($sd);
		RAIDHotRemove($md, $sd) if (IsFailed($sd));
		RAIDHotAdd($md, $sd);
	}
	if (Failures() == 0) {
		Exit(0);
	} else {
		print("Unable to resolve all failures\n");
		ShowFailed();
	}
}

sub Members {
	# Takes an md name and returns a list of sd's that 
	# are its devices
	my @members;
	my $md = shift();
	foreach my $sd (keys(%{$mdConfig{$md}})) {
		push(@members, $sd);
	}
	return(@members);
}

sub RAIDSetFaulty {
	my ($md, $sd) = @_;
	system("$MDADM $md -f $sd 2> /dev/null 1> /dev/null");
}

sub RAIDHotRemove {
	my ($md, $sd) = @_;
	system("$MDADM $md -r $sd 2> /dev/null 1> /dev/null");
}

sub RAIDHotAdd {
	my ($md, $sd) = @_;
	system("$MDADM $md -a $sd 2> /dev/null 1> /dev/null");
}

sub IsUp {
	my $sd = shift();
	my ($md) = MemberOf($sd);
	if (defined($mdState{$md}{$sd}) &&
	    $mdState{$md}{$sd} eq 'U') {
		return(1);
	}
	return(0);
}

sub IsDown {
	my $sd = shift();
	my ($md) = MemberOf($sd);
	if (!defined($mdState{$md}{$sd})) {
		return(1);
	}
	return(0);
}

sub IsFailed {
	my $sd = shift();
	my ($md) = MemberOf($sd);
	if (defined($mdState{$md}{$sd}) && 
	    $mdState{$md}{$sd} eq 'F') {
	    	return(1);
	}
	return(0);
}

sub MemberOf {
	# Takes an sd argument and returns the md it's a member of
	# If given an sd without partition (e.g., sda vs. sda1)
	# returns all md's the sd is a member of
	my @list;
	my $arg = shift();
	foreach my $md (keys(%mdConfig)) {
		foreach my $sd (keys(%{$mdConfig{$md}})) {
			if ($sd =~ m/$arg/) {
				push(@list, $md);
			}
		}
	}
	return(@list);
}

sub IsPresentSCSI {
	# Report presence of sd in /proc/scsi/scsi
	# ARGV STANDARD
	my $arg = shift();
	# Don't iterate through %scsi, we have the hash key - ARGV STANDARD
	foreach my $drive (keys(%scsi)) {
		if ($drive =~ m/$arg/) {
			return(1);
		}
	}
	return(0);
}

sub IsPresentRAID {
	# Report presence of sd, up or down, in any md
	# ARGV STANDARD
	my $arg = shift();
	foreach my $md (keys(%mdState)) {
		foreach my $sd (keys(%{$mdState{$md}})) {
			if ($sd =~ m/$arg/) {
				return(1);
			}
		}
	}
	return(0);
}

sub ReinsertDrive {
	# Reinsert drive in /proc/scsi/scsi and repartition drive using 
	# data stored by SaveDrive()
	my ($bus, $chan, $id, $lun);
	my $drive = shift();

	if (-r "$TMPSCSI$drive") {
		open(SCSI, "<", "$TMPSCSI$drive");
		($bus, $chan, $id, $lun) = split(/\s+/, <SCSI>);
		close(SCSI);
		unlink("$TMPSCSI$drive");
		qx(echo "scsi add-single-device $bus $chan $id $lun" > /proc/scsi/scsi);
		ReadSCSI();
		if (!IsPresentSCSI($drive)) {
			Exit(2, "Inserting $drive into SCSI chain failed");
		}
	} else { 
		Exit(2, "No information on how to reinsert /dev/$drive");
	}

	if (-r "$TMPSFDISK$drive") {
		if (system("$SFDISK /dev/$drive < $TMPSFDISK$drive " .
		    "2> /dev/null 1> /dev/null")) {
			Exit(1, "Failed to partition /dev/$drive");
		} else {
			unlink("$TMPSFDISK$drive");
		}
	} else {
		Exit(2, "No information on how to repartition /dev/$drive");
	}
}
	
sub SanityCheck {
	if (! -r $CONFIG || ! -r $STATE || ! -r $SCSI || ! -x $GRUB ||
	    ! -x $SCSIINFO || ! -x $SFDISK || ! -x $MDADM) {
		Exit(2, "Missing a required file or utility");
	}
}

sub SaveDrive {
	# Store partition table and SCSI path
	my $drive = shift();

	if (system("$SFDISK -d /dev/$drive 1> $TMPSFDISK$drive")) {
		unlink("$TMPSFDISK$drive");
		Exit(2, "Failed to dump existing partition map from /dev/$drive");
	}

	open(SCSI, ">", "$TMPSCSI$drive") or 
	    die("Couldn't open $TMPSCSI$drive for writing");
	print(SCSI $scsi{"/dev/$drive"});
	close(SCSI);
}

sub Grub {
	# Write grub boot labels to anything defined as a boot drive
	for (@BOOTDRIVES) {
		if (IsPresentSCSI($_)) {
			open(GRUB, ">", "$TMPGRUB$_") or 
			    die("Couldn't open $TMPGRUB$_ for writing");
			print(GRUB "device (hd0) /dev/$_\n" .
			    "root (hd0,0)\n" .
			    "setup (hd0)\n");
			close(GRUB);
			system("$GRUB < $TMPGRUB$_ 2> /dev/null 1> /dev/null");
			unlink("$TMPGRUB$_");
		}
	}
}

sub Exit {
	my ($exitVal, $errMsg) = @_;
	print("$errMsg\n") if defined($errMsg);
	exit($exitVal);
}