#!/usr/bin/perl -w

# ======================================================================

# Copyright (c) 2013 Kenneth C. Schalk

# A script which may, under the right circumstances, help to rescue a
# broken Linux software RAID using device mapper snapshots.

# You may redistribute and/or modify this script under the terms of
# the GNU General Public License as published by the Free Software
# Foundation:

# http://www.gnu.org/licenses/gpl.html

# This script is provided for educational purposes only with no
# warranty express or implied.  It should not be considered safe or
# robust.  It will most likely not work for you without modifications
# to suit your particular situation.  Since it must be run as root it
# has the potential to destroy your system and you should obviously
# not trust it.  If you don't feel comfortable reading, auditing, and
# modifying this script DO NOT USE IT.

# If you have several hard drive images from a failed RAID array in files named:

#   sda1.img sdb3.img sdd3.img sde3.img

# Then you might run this script like this:

#   sudo recreate-raid-via-dm-snapshot.pl --start sda1.img sdb3.img sdd3.img sde3.img

# Later you might reverse that like this:

#   sudo recreate-raid-via-dm-snapshot.pl --stop sda1.img sdb3.img sdd3.img sde3.img

# For more information, see:

# http://st.xorian.net/blog/2013/03/using-the-linux-device-mapper-to-rescue-a-failed-raid

# ======================================================================

use POSIX;
use Getopt::Long;
use File::Basename;

%Options = (
	    # Prefix to use for created device mapper files
	    'dm-prefix' => 'rr-',

	    # Suffix to use for created CoW files
	    'cow-suffix' => '-cow',

	    # How big should the CoW files be?  (They're created as
	    # sparse files, but this size limits how much change they
	    # can hold.)
	    'cow-size' => '1G',

	    # Where should the CoW files be created?
	    'cow-dir' => '.',

	    # Should the CoW files be deleted?  (If not, existing ones
	    # may be re-used.)
	    'delete-cow' => 0,

	    # Should we be starting (setting up) or stopping (tearing
	    # down)?
	    'start' => 0,
	    'stop' => 0,

	    # Should we print lots of extra information about what
	    # we're doing?
	    'debug' => 0,
	    );

GetOptions(\%Options,
	   'dm-prefix=s',
	   'cow-suffix=s',
	   'cow-size=s',
	   'cow-dir=s',
	   'start', 'stop', 'debug');

# Get the size of a file in bytes

sub file_bytes
  {
    my $fname = shift;
    my ($dev,$ino,$mode,$nlink,$uid,$gid,$rdev,$size,
	$atime,$mtime,$ctime,$blksize,$blocks)
      = stat($fname);
    return $size;
  }

# Get information about an existing md component

sub md_info
  {
    my $fname = shift;
    my %result = ();
    unless(open MDADM_EXAMINE, "/sbin/mdadm --examine $fname |")
      {
	return undef;
      }
    while(<MDADM_EXAMINE>)
      {
	chomp;
	if(/^\s*Version\s*:\s*(\S+)/)
	  {
	    $result{'version'} = $1;
	    $result{'version'} =~ s/\.00$//;
	  }
	elsif(/^\s*Raid Level\s*:\s*raid(\d+)/)
	  {
	    $result{'raid_level'} = $1;
	  }
	elsif(/^\s*Raid Devices\s*:\s*(\d+)/)
	  {
	    $result{'raid_device_count'} = $1;
	  }
	elsif(/^\s*(Array )?UUID\s*:\s*(\S+)/)
	  {
	    $result{'uuid'} = $2;
	  }
	elsif(/^\s*this\s+(\d+)/)
	  {
	    $result{'number'} = $1;
	  }
	elsif(/^\s*Device Role\s*:.*device (\d+)/)
	  {
	    $result{'number'} = $1;
	  }
	elsif(/^\s*Chunk Size\s*:\s*(\d+)K/)
	  {
	    $result{'chunk_size'} = $1;
	  }
      }
    unless(close MDADM_EXAMINE)
      {
        if($!)
	  {
            print STDERR ("Error closing command pipe for 'mdadm --examine $fname': $!") unless($! =~ /No child processes/i);
	  }
        else
	  {
            print STDERR ("Exit status $? from 'mdadm --examine $fname'");
	  }
      }
    return undef if(scalar(keys(%result)) < 1);
    $result{'blocks'} = POSIX::ceil(file_bytes($fname) / 512);
    return \%result;
  }

# Compare specific keys in two hashes for equality.  Return a list of
# the keys which don't match.

sub compare_keys
  {
    my $hash1 = shift;
    my $hash2 = shift;
    my @result = ();

    foreach my $key (@_)
      {
	if($hash1->{$key} ne $hash2->{$key})
	  {
	    push @result, $key;
	  }
      }

    return @result;
  }

# Look for an existing loop device for a given file

sub existing_loop_dev
  {
    my $fname = shift;
    my $loop_dev = `/sbin/losetup -j $fname`;
    if(($? == 0) &&
       ($loop_dev =~ m|^(/dev/loop\d+):.*|))
      {
	return $1;
      }
    return undef;
  }

# Set up a loop device for a file, or return the existing one if it
# already exists

sub start_loop_dev
  {
    my $fname = shift;
    my $loop_dev = existing_loop_dev($fname);
    if(!defined($loop_dev))
      {
	$loop_dev = `/sbin/losetup -f --show $fname`;
	if($? == 0)
	  {
	    chomp $loop_dev;
	    print("Started loop $loop_dev for $fname\n") if($Options{'debug'});
	  }
	else
	  {
	    $loop_dev = undef;
	  }
      }
    elsif($Options{'debug'})
      {
	print("($loop_dev already exists for $fname)\n");
      }

    return $loop_dev;
  }

# Stop the loopback device for a file, if it exists

sub stop_loop_dev
  {
    my $fname = shift;
    my $loop_dev = existing_loop_dev($fname);
    if(defined($loop_dev))
      {
	if(system("/sbin/losetup -d $loop_dev") != 0)
	  {
	    print STDERR ("Warning: couldn't stop loop device ", $loop_dev, " for ", $fname, "\n");
	    return 0;
	  }
	elsif($Options{'debug'})
	  {
	    print("(Stopped $loop_dev for $fname)\n");
	  }
      }
    elsif($Options{'debug'})
      {
	print("(No loop device found for $fname)\n");
      }
    return 1;
  }

# Convert an original path into the name of the CoW file to be used
# with it

sub cow_fname
  {
    my $orig_name = shift;
    return $Options{'cow-dir'}."/".basename($orig_name).$Options{'cow-suffix'};
  }

# Convert an original path into a name we'll use with the device
# mapper.

sub dm_name
  {
    my $orig_name = shift;
    return $Options{'dm-prefix'}.basename($orig_name);
  }

# Create a CoW file if it doesn't exist already

sub make_cow_file
  {
    my $cow_fname = shift;
    if(!-f $cow_fname)
      {
	my $cmd = "dd if=/dev/zero of=$cow_fname bs=1 seek=".$Options{'cow-size'}." count=1";
	print("Creating CoW file: $cmd\n") if($Options{'debug'});
	unless(system($cmd) == 0)
	  {
	    return 0;
	  }
      }
    return 1;
  }

# Call dmsetup to create a device if needed

sub dm_setup
{
    my $dm_name = shift;
    my $table_str = shift;

    my $dm_status = `/sbin/dmsetup status $dm_name 2>&1`;
    if($? == 0)
    {
	chomp $dm_status;
	my @dm_status_bits = split ' ', $dm_status;
	my @table_bits = split ' ', $table_str;
	if(($dm_status_bits[0] eq $table_bits[0]) &&
	   ($dm_status_bits[1] eq $table_bits[1]) &&
	   ($dm_status_bits[2] eq $table_bits[2]))
	{
	    return $dm_name;
	}
	else
	{
	    print STDERR ($dm_name, " (", $table_bits[2], "): exists but not correct!\n");
	}
    }
    else
    {
	open DMSETUP, "|/sbin/dmsetup create $dm_name";
	print DMSETUP ($table_str, "\n");
	close DMSETUP;
	if($? == 0)
	{
	    return $dm_name;
	}
	else
	{
	    print STDERR ($dm_name, ": error creating!\n");
	}
    }
    return undef;
}

# Call dmsetup to remove a device.

sub dm_remove
  {
    my $dm_name = shift;
    return if(!-e ("/dev/mapper/".$dm_name));
    unless(system("/sbin/dmsetup remove $dm_name") == 0)
      {
	print STDERR ($dm_name, ": couldn't de-activate!\n");
	return 0;
      }

    return 1;
  }

# Erase the software RAID superblock on a device.

sub md_zero_superblock
  {
    my $dev = shift;
    my $cmd = "mdadm --zero-superblock ".$dev;
    print("Zeroing md superblock: $cmd\n") if($Options{'debug'});
    unless(system($cmd) == 0)
      {
	return 0;
      }
    return 1;
  }

# Information about all the input files
%finfo = ();

# The first one we successfully get information about
my $compare_info = undef;

# Examine all the arguments first
foreach my $fname (@ARGV)
    {
      my $info = md_info($fname);
      next if(!defined($info));
      if($Options{'debug'})
	{
	  print($fname, ":\n");
	  foreach my $key (sort(keys(%$info)))
	    {
	      print("\t", $key, " = ", $info->{$key}, "\n");
	    }
	}
      if(!defined($compare_info))
	{
	  $compare_info = $info;
	}
      else
	{
	  my @mismatches = compare_keys($info, $compare_info,
					'uuid', 'version', 'raid_level',
					'raid_device_count', 'chunk_size',
					'blocks');
	  if(scalar(@mismatches) > 0)
	    {
	      print STDERR ("Rejecting $fname : mismatch on ", join(",", @mismatches), "\n");
	      next;
	    }

	  if($info->{'number'} >= $info->{'raid_device_count'})
	    {
	      print STDERR ("Rejecting $fname : bad device number (",
			    $info->{'number'}, " in an array of ",
			    $info->{'raid_device_count'},
			    " devices)\n");
	      next;
	    }
	}
      $finfo{$fname} = $info;
    }

if($Options{'start'})
  {
    # Make sure needed kernel modules are loaded.
    unless(system("/sbin/modprobe loop max_loop=64") == 0)
      {
	print STDERR ("Couldn't load loop module\n");
	exit(1);
      }
    unless(system("/sbin/modprobe dm-snapshot") == 0)
      {
	print STDERR ("Couldn't load dm-snapshot module\n");
	exit(1);
      }

    if(!defined($compare_info))
      {
	print STDERR ("No usable devices found\n");
      }
    elsif(($compare_info->{'raid_level'} == 5) &&
	  (($compare_info->{'raid_device_count'} - 1) >
	   scalar(keys(%finfo))))
      {
	print STDERR ("Not enough usable devices found the re-assemble RAID5 (",
		      scalar(keys(%finfo)), " found, ",
		      ($compare_info->{'raid_device_count'} - 1), " needed)\n");
      }
    elsif($compare_info->{'raid_level'} == 5)
      {
	my @assemble_order = ("missing") x $compare_info->{'raid_device_count'};
	foreach my $fname (keys(%finfo))
	  {
	    my $info = $finfo{$fname};
	    $info->{'cow_fname'} = cow_fname($fname);
	    unlink $info->{'cow_fname'} if($Options{'delete-cow'} && -f $info->{'cow_fname'});
	    if(!make_cow_file($info->{'cow_fname'}))
	      {
		print STDERR ("Couldn't create Cow file ", $info->{'cow_fname'},
			      " for ", $fname, "\n");
		exit(1);
	      }
	    $info->{'cow_loop'} = start_loop_dev($info->{'cow_fname'});
	    if(!defined($info->{'cow_loop'}))
	      {
		print STDERR ("Couldn't start loopback for CoW file file ", $info->{'cow_fname'},
			      " for ", $fname, "\n");
		exit(1);
	      }
	    $info->{'dev'} = start_loop_dev($fname);
	    if(!defined($info->{'dev'}))
	      {
		print STDERR ("Couldn't start loopback for ", $fname, "\n");
		exit(1);
	      }
	    $info->{'origin_name'} = dm_setup(dm_name($fname),
					      ("0 ".$info->{'blocks'}." snapshot-origin ".$info->{'dev'}));
	    if(!defined($info->{'origin_name'}))
	      {
		print STDERR ("Couldn't start snapshot-origin for ", $fname, "\n");
		exit(1);
	      }

	    $info->{'snapshot_name'} = dm_setup(dm_name($info->{'cow_fname'}),
						("0 ".$info->{'blocks'}." snapshot ".$info->{'dev'}." ".$info->{'cow_loop'}." p 128"));
	    if(!defined($info->{'snapshot_name'}))
	      {
		print STDERR ("Couldn't start snapshot for ", $fname, "\n");
		exit(1);
	      }

	    $info->{'snapshot_dev'} = "/dev/mapper/".$info->{'snapshot_name'};
	    md_zero_superblock($info->{'snapshot_dev'});

	    $assemble_order[$info->{'number'}] = $info->{'snapshot_dev'};
	  }

	if($Options{'debug'})
	  {
	    print("Proposed RAID5 assembly order: ",
		  join(" ", @assemble_order), "\n");
	  }

	my $cmd = ("mdadm --create /dev/md/rr --assume-clean -l ".$compare_info->{'raid_level'}.
		   " -n ".$compare_info->{'raid_device_count'}.
		   " -c ".$compare_info->{'chunk_size'}.
		   " --metadata=".$compare_info->{'version'}." ".
		   join(" ", @assemble_order));
	print("Creating array: $cmd\n") if($Options{'debug'});
	unless(system($cmd) == 0)
	  {
	    print STDERR ("Error creating array\n");
	    exit(1);
	  }
      }
    else
      {
	print STDERR ("Sorry, the script doesn't seem to support your situation (RAID level other than 5?)\n");
      }
  }
elsif($Options{'stop'})
  {
    # Note, this doesn't stop the md device.  You'll have to do that
    # manually first.  (And if it's being used by another subsystem
    # like LVM, you may need to take other steps first.)

    foreach my $fname (keys(%finfo))
      {
	my $info = $finfo{$fname};
	# Stop the snapshot-origin device
	dm_remove(dm_name($fname));
	$info->{'cow_fname'} = cow_fname($fname);
	# Stop the snapshot device
	dm_remove(dm_name($info->{'cow_fname'}));
	if(-f $info->{'cow_fname'})
	  {
	    if(stop_loop_dev($info->{'cow_fname'}) &&
	       $Options{'delete-cow'})
	      {
		unlink $info->{'cow_fname'}
	      }
	  }
	stop_loop_dev($fname);
      }
  }