#!/usr/bin/perl -w

# This is a perl script for cloning disk to another one,
# This can run on SystemRescue (tested Ver11.00).

use Data::Dumper;
use Getopt::Long;
use JSON::PP qw(decode_json);
use Clone qw(clone);
use Fcntl;
use File::Temp qw(tempfile);
use strict;
use warnings;

# set $ENV{'PATH'} for security
$ENV{'PATH'} = "/usr/bin:/bin";

# Global Values
our %opts;
our $srci;
our $dsti;

# align sector into 1MB
our $ALIGNSIZE = 1024*1024; 
# first partition start from 1MB
our $FIRSTPARTSTART = 1024*1024;
# reserved size(2MB) of the last of the disk, as Windows
our $LASTRESERVE = 2*1024*1024;

# print usage
sub print_usage {
	print "clone disk from src_device to dst_device\n";
	print "usage: $0 [options] <src_device> <dst_device>\n";
	print "   ex. $0 /dev/sda /dev/sdb\n";
	print "options:\n";
	print "--help     : show this message\n";
	print "--test     : test only, nothing changed\n";
	print "--compact  : shrink partitions to the top of the destination disk\n";
	print "--winlocate: re-locate partitions as of Windows default\n";
	print "--reorder  : re-number the partitions from the top of the disk\n";
	print "--winonly  : clone Windows partitions only\n";
	print "--verbose  : display further information\n";
	print "--partonly : partitioning only, not copy filesystems\n";
	print "--forcedd  : not use partclone but use dd for all partitions\n";
	print "--only=<part>    : clone only specified partition(s), delimited by ','\n";
	print "--except=<part>  : clone except specified partition(s), delimited by ','\n";
	print "--partdata=<file>: read partition data from file insted of <src_device>\n";
	exit 1;
}


# run command and return the output as array
# Argument: command string, error message
# Note: $? can be checked after this function
# when error($? != 0):
#   if $errmes is undefined, return all output
#   if $errmes is defined, print command and its output then die
sub cmd {
	my ($cmdline, $errmes) = @_;
	if ($cmdline =~ /[;\n|]/) {
		print "# $cmdline\n";
		die "command includes invalid charactors";
	}
	my @ret = `$cmdline 2>&1`;
	my $output = "# $cmdline\n".join('', @ret);
	print $output if ($opts{verbose});
	return @ret if ($? == 0 || !defined($errmes));

	print $output if (!$opts{verbose});
	die $errmes;
}
# run command and return the output as string
# Argument: command string, error message
sub cmd1 {
	my ($cmdline, $errmes) = @_;
	my $ret = join('', cmd($cmdline, $errmes));
	chomp $ret;
	return $ret;
}


# get partition device(ex. /dev/nvme0n1p1, /dev/sda1)
# Argument: whole of device(ex. /dev/sda), partition(ex./dev/sda1, or 1)
sub get_partdev {
	my ($dev, $part) = @_;
	$part =~ s/.*?(\d+)$/$1/;
	return ($dev =~ /\d$/) ? "${dev}p$part" : "$dev$part";
}


# align partition sector *after* the argument
# Argument: sector
# Return  : aligned start sector
sub align_sector {
	my ($sect, $sectsiz) = @_;
	my $elm = int($ALIGNSIZE/$sectsiz);
	return int(($sect+$elm-1)/$elm)*$elm;
}
# align partition sector *before* the argument
sub align_sector_before {
	my ($sect, $sectsiz) = @_;
	my $elm = int($ALIGNSIZE/$sectsiz);
	return int($sect/$elm)*$elm;
}
# convert sector number for the disk($from_sectsize) to one($to_sectsize)
sub convert_sector {
	my ($sect, $from_sectsiz, $to_sectsiz) = @_;
	return int((($sect*$from_sectsiz)+$to_sectsiz-1)/$to_sectsiz);
	# do not align, for DOS pt copy.
}
# size to sector/sector
sub size_to_sector {
	my ($siz, $sectsiz) = @_;
	return int(($siz+$sectsiz-1)/$sectsiz);
}
# get whole disk size(in byte)
sub get_disk_size {
	return cmd1("blockdev --getsize64 $_[0]", "Fail to get size");
}
# get logical sector size(in byte, ex. 512, 4096)
sub get_sector_size {
	return cmd1("blockdev --getss $_[0]", "Fail to get size");
}


# check if it is windows partition on gpt
sub is_windows_partition {
	my $t = $_[0];
	return
		$t eq "C12A7328-F81F-11D2-BA4B-00A0C93EC93B" || # EFI system partition
		$t eq "E3C9E316-0B5C-4DB8-817D-F92DF00215AE" || # MS reserved part
		$t eq "EBD0A0A2-B9E5-4433-87C0-68B6B72699C7" || # Data partition
		$t eq "5808C8AA-7E8F-42E0-85D2-E1E90434CFB3" || # LDM metadata partition
		$t eq "AF9B60A0-1431-4F62-BC68-3311714A69AD" || # LDM data partition
		$t eq "DE94BBA4-06D1-4D40-A16A-BFD50179D6AC";  # Recovery partition
}


# check if the partition table is gpt/dos.
sub is_gpt { return $_[0]->{label} eq "gpt"; }
sub is_dos { return $_[0]->{label} eq "dos"; }


# get disk information by "sfdisk -J $dev" and return reference of data
# Argument: whole device(ex. /dev/sda)
# Return  : reference of the disk information
sub get_disk_info {
	print "Getting source device information...\n";
	my $dev = $_[0];
	my @lines = cmd("sfdisk -J $dev", "Cannot get device info");
	my $json = (decode_json(join('', @lines)))->{partitiontable};
	# get disksize(internal data, sector number of the disk)
	# --getsz shows "size in 512byte/sector", then 4096byte/sector?
	my $size = get_disk_size($dev);
	$json->{disksize} = $size/$json->{sectorsize}; # must be integer
	return $json;
}


# Get sfdisk string of the disk
# Argument: reference of the disk information
# Return  : string used for sfdisk
sub get_sfdisk_string {
	my $pt = $_[0];
	my $ret = "";

	$ret .= "label: $pt->{label}\n";
	$ret .= "label-id: $pt->{id}\n";
	$ret .= "device: $pt->{device}\n";
	$ret .= "unit: $pt->{unit}\n";
	# "if (...)" is necessary to apply both gpt/dos partition table.
	$ret .= "first-lba: $pt->{firstlba}\n" if ($pt->{firstlba});
	$ret .= "last-lba: $pt->{lastlba}\n"   if ($pt->{lastlba});
	$ret .= "sector-size: $pt->{sectorsize}\n\n";
	foreach (@{$pt->{partitions}}) {
		$ret .= "$_->{node} : start=$_->{start}, size=$_->{size}, type=$_->{type}";
		$ret .= ", uuid=$_->{uuid}"       if ($_->{uuid});
		$ret .= ", name=\"$_->{name}\""   if ($_->{name});
		$ret .= ", attrs=\"$_->{attrs}\"" if ($_->{attrs});
		$ret .= ", bootable"              if ($_->{bootable});
		$ret .= "\n";
	}
	return $ret;
}


# Copy MBR from src to dst device(only 1MB... good enough for most cases)
# Argument: src_device, dst_device
# Return  : nothing
sub copy_mbr {
	print "Copying MBR...\n";
	my ($src, $dst) = @_;
	(-b $src) || die "$src is not block device file";
	(-b $dst) || die "$dst is not block device file";
	cmd("dd if=$src of=$dst bs=1M count=1", "Fail to copy MBR");
}


# Apply disk information by running "sfdisk $dev < tmpfile"
# Argument: reference of the disk information
sub apply_disk_info {
	print "Applying destination device information...\n";
	my $pt = $_[0];
	my ($fh, $fnam) = tempfile();
	print $fh get_sfdisk_string($pt);
	close($fh);
	# it is important to make tempfile to check it after any errors.
	cmd("sfdisk $pt->{device} < $fnam", "Fail to apply partdata");
	unlink $fnam;
	cmd("blockdev --rereadpt $pt->{device}", "Fail to re-read partition table");

	# "blockdev --rereadpt" could refresh /dev/*, so wait for the device files
	foreach (@{$pt->{partitions}}) {
		for (my $i= 0; $i < 3 && (!-b $_->{node}); $i++) {
			sleep(1);
		}
		(-b $_->{node}) || die "$_->{node} is not block device file";
	}
}


# change the device name in the disk information
# Argument: reference of the disk information, device
sub change_device_name {
	print "Configuring destination device...\n";
	my ($pt, $dev) = @_;
	$pt->{device} = $dev;

	my $fromsiz = $pt->{sectorsize};
	my $tosiz = get_sector_size($dev);
	$pt->{sectorsize} = $tosiz;
	$pt->{disksize} = get_disk_size($dev)/$tosiz; # must be integer

	# This is gpt only, "dos" does not have {lastlba};
	$pt->{lastlba} = $pt->{disksize}-34 if (is_gpt($pt));

	foreach (@{$pt->{partitions}}) {
		$_->{node}  = get_partdev($dev, $_->{node});
		$_->{start} = convert_sector($_->{start}, $fromsiz, $tosiz);
		$_->{size}  = convert_sector($_->{size},  $fromsiz, $tosiz);
	}
}


# process "--only" and "--except" options
# Argument: reference of the disk information
sub process_only_except_option {
	my @pts = @{$_[0]->{partitions}};
	if (defined($opts{only})) {
		my @devs = split(/,/, $opts{only});
		@pts = grep { my $d = $_->{node};  grep(/^$d$/, @devs) } @pts;
	}
	if (defined($opts{except})) {
		my @devs = split(/,/, $opts{except});
		@pts = grep { my $d = $_->{node}; !grep(/^$d$/, @devs) } @pts;
	}
	$_[0]->{partitions} = \@pts;
}


# Comact/shrink partitions, means place all partitions without unused space
# Argument: reference of the disk information
sub compact_partitions {
	print "Compacting partitions...\n";
	my @pts = @{$_[0]->{partitions}};
	my $sectsiz = $_[0]->{sectorsize};

	# 1st partition should be after 1MB(2048s on 512b/sector)
	my $start = size_to_sector($FIRSTPARTSTART, $sectsiz);
	# place all partitions without gap
	foreach my $t (@pts) {
		$t->{start} = align_sector($start, $sectsiz);
		$start = $t->{start} + $t->{size};
	}
}


# locate partitions as Windows:
#     try to locate EFI first, reserved next, then C and Recovery last
# Argument: reference of the disk information
sub winlocate_partitions {
	print "Relocating partitions for Windows...\n";
	my $pt = $_[0];
	my @pts = @{$pt->{partitions}};
	my $sectsiz = $pt->{sectorsize};

	# search Recovery partition and place it on the last of the disk
	my @recov = grep { $_->{type} eq "DE94BBA4-06D1-4D40-A16A-BFD50179D6AC"} @pts;
	my $end = $pt->{disksize}-1-size_to_sector($LASTRESERVE, $sectsiz);
	# as Windows GTP, last 2MB(4096 with 512b/sect) is reserved.
	foreach my $t (@recov) {
		$t->{start} = align_sector_before($end-$t->{size}+1, $sectsiz);
		$end = $t->{start}-1;
		@pts = grep($_ ne $t, @pts); # remove the partition from @pts
		# I know there are not two or more recover partitions...
	}

	my $start = size_to_sector($FIRSTPARTSTART, $sectsiz);
	# search "EFI system partition" and place it on the top of the disk
	my @efis = grep {$_->{type} eq "C12A7328-F81F-11D2-BA4B-00A0C93EC93B" } @pts;
	foreach my $t (@efis) {
		$t->{start} = align_sector($start, $sectsiz);
		$start = $t->{start} + $t->{size};
		@pts = grep($_ ne $t, @pts); # remove the partition from @pts
		# I know there are not two or more efi partitions...
	}

	# search "Microsoft reserved partition" and place it next to EFI
	my @reserv = grep {$_->{type} eq "E3C9E316-0B5C-4DB8-817D-F92DF00215AE"} @pts;
	foreach my $t (@reserv) {
		$t->{start} = align_sector($start, $sectsiz);
		$start = $t->{start} + $t->{size};
		@pts = grep($_ ne $t, @pts); # remove the partition from @pts
		# I know there are not two or more reserved partitions...
	}

	# place all other partitions to the next
	foreach my $t (@pts) {
		$t->{start} = align_sector($start, $sectsiz);
		$start = $t->{start} + $t->{size};
	}
}


# modify partition order as "starting earlier is first"
# Argument: reference of the disk information
sub reorder_partitions {
	print "Reordering partitions...\n";
	my $pt = $_[0];
	my $ptsp = $pt->{partitions};

	@$ptsp = sort { $a->{start} <=> $b->{start} } @$ptsp;
	for(my $i = 0; $i <= $#$ptsp; $i++) {
		$ptsp->[$i]{node} = get_partdev($pt->{device}, $i+1);
	}
}


# get filesystem name by fsstat and trying mount 
# Argument: partition(ex. /dev/sda1, /dev/sdb1)
# return  : filesystem name or "" if nothing found
sub get_filesystem_name {
	my ($part) = @_;
	my $ret = cmd1("fsstat -t $part");
	return $ret if ($? == 0 && $ret ne "");
	# Note: parted is not suitable, as it refreshs /dev/*, often lose the devfile
	# So trying mouinting it and check the filesystem here
	cmd("mount -r $part /mnt");
	return "" if ($? != 0);
	my @partdat = cmd("mount");
	cmd("umount $part", "Fail to umount");
	@partdat = grep(/^$part on /, @partdat);
	return ($#partdat != 0) ? "" : (split(/\s+/, $partdat[0]))[4];
}


# Copy/clone partition
# Argument: src_device, dst_deviceZ
sub copy_partition {
	my($src, $dst) = @_;

	(-b $src) || die "$src is not block device file";
	(-b $dst) || die "$dst is not block device file";
	my $fs = ($opts{forcedd}) ? "" : get_filesystem_name($src);
	if ($fs ne "") {
		print "    Copying $fs from $src to $dst...\n";
		# partclone buffer size(-z) is 100MB
		cmd("partclone.$fs -z 104857600 -b -s $src -o $dst");
		return if ($? == 0);
		print "      Fail to run partclone.$fs, try dd...\n";
		# fail safe, go through down
	}
	# for any other partitions
	print "    Cloning partition from $src to $dst...\n";
	cmd("dd if=$src of=$dst bs=100M", "Fail to clone partition");
}


# copy all partitions in the dst device from src to dst for gpt
# Argument: src_device_information, dst_device_information
sub copy_all_partitions_gpt {
	print "Copying partitions on gpt...\n";
	my ($srci, $dsti) = @_;
	my $srcpt = $srci->{partitions};
	my $dstpt = $dsti->{partitions};

	# sort by partition size, not to spend time to wait errors occur
	my @dstpary = sort { $a->{size} <=> $b->{size} } @$dstpt;

	foreach my $dstp (@dstpary) {
		my @parts = grep {$dstp->{uuid} eq $_->{uuid}} @$srcpt;
		($#parts == 0) || die "Wrong src, corresponding to $dstp->{node}";
		my $srcp = $parts[0];
		copy_partition($srcp->{node}, $dstp->{node});
	}
}


# copy all partitions from src to dst for dos pt
# Argument: src_device_information, dst_device_information
sub copy_all_partitions_dos {
	print "Copying partitions on dos...\n";
	my ($srci, $dsti) = @_;
	my $srcpt = $srci->{partitions};
	my $dstpt = $dsti->{partitions};

	# dos has many restrictions...
	($#$srcpt == $#$dstpt) || die "src/dst partition number is different";

	for(my $i=0; $i <= $#$srcpt; $i++) {
		my ($srcp, $dstp) = ($srcpt->[$i], $dstpt->[$i]);
		next if ($srcp->{type} eq "f"); # skip extend partition
		copy_partition($srcp->{node}, $dstp->{node});
	}
}


# read sfdisk data which is generated by "sfdisk -d <dev>"
# Argument: filename of the data(which is genereted by "sfdisk -d <dev>"
# Return  : reference of the hash of the sfdisk data
sub read_sfdisk_data {
	my $file = $_[0];
	my $ret = {};

	open(my $fh, "< $file") || die "Fail to open file($file)";
	while (<$fh>) {
		chomp;
		last if (/^\s*$/);
		my ($key, $val) = split(/\s*:\s+/, $_);
		$key =~ s/label-id/id/;
		$key =~ s/-//g; # remove '-' like "first-lba" to "firstlba"
		$key =~ s/^\s*(.*?)\s*$/$1/;
		$val =~ s/^\s*(.*?)\s*$/$1/;
		$ret->{$key} = $val;
	}
	$ret->{partitions} = [];
	for (my $i = 0; <$fh>; $i++) {
		chomp;
		my ($key, $val) = split(/\s*:\s+/, $_);
		$key =~ s/^\s*(.*?)\s*$/$1/;
		$val =~ s/^\s*(.*?)\s*$/$1/;
		$ret->{partitions}[$i]->{node} = $key;
		foreach (split(/\s*,\s*/, $val)) {
			# FIXME: split() can't work correctly if there's ',' in string
			# ex: 'a=b, c="d,e,f", f=g'
			if (/^bootable$/) {
				$ret->{partitions}[$i]->{bootable} = 1;
			} else {
				my ($key, $val) = split(/\s*=\s*/);
				$val =~ s/^"(.*)"$/$1/;
				$ret->{partitions}[$i]->{$key} = $val;
			}
		}
	}
	close($fh);
	return $ret;
}


# Ask Yes or No and return true or false
# Argument: message
sub ask_yn {
	my $msg = $_[0];
	while(1) {
		print "$msg [y/n] ";
		my $line = <STDIN>;
		chomp($line);
		($line eq 'y') && return 1;
		($line eq 'n') && return 0;
		print "Please input \"y\" or \"n\"\n";
	}
}


# remove non-windows partitions
# Argument: reference of the disk information, device
sub remove_nonwin_partitions {
	my $pt = $_[0];
	my @ptary = @{$pt->{partitions}};
	@ptary = grep { is_windows_partition($_->{type}) } @ptary;
	$pt->{partitions} = \@ptary;
}


# main function
{ 
	print "Disk cloner ver 2024-05-03-0\n";

	my $oret = GetOptions(\%opts, "help", "test", "compact", "winlocate", 
				"reorder", "winonly", "verbose", "partdata=s",
				"forcedd", "partonly", "only=s", "except=s");
	($oret && $#ARGV == 2-1 && !$opts{help}) || print_usage();
	my ($srcdev, $dstdev) = @ARGV;
	
	$srci = get_disk_info($srcdev);

	# label should be "gpt" or "dos"
	if (!is_gpt($srci) && !is_dos($srci)) {
		die "disk($srci->{device}) label($srci->{label}) must be gpt/dos";
	}
	if (is_dos($srci) &&
		($opts{compact} || $opts{winlocate} || $opts{reorder} || $opts{winonly} ||
		 $opts{except} || $opts{only})) {
		die "comact/winlocate/reorder/winonly/except/only can not be used for dos";
	}

	if ($opts{partdata}) {
		$dsti = read_sfdisk_data($opts{partdata});
		change_device_name($dsti, $dstdev);
	} else {
		$dsti = clone($srci);
		process_only_except_option($dsti);
		change_device_name($dsti, $dstdev);
		($opts{winonly})   && remove_nonwin_partitions($dsti);
		($opts{compact})   && compact_partitions($dsti);
		($opts{winlocate}) && winlocate_partitions($dsti);
		($opts{reorder})   && reorder_partitions($dsti);
	}

	my $sfdiskstr = get_sfdisk_string($dsti);
	print "\nDestination disk($dsti->{device}) will be changed to:\n";
	print "$sfdiskstr\n";
	if ($opts{test}) {
		print "Do nothing, due to test\n";
		exit 0;
	}

	if (!ask_yn( "ALL PARTITIONS on $dsti->{device} will be REMOVEd. OK?")) {
		die "aborted";
	}

	# copy mbr for OS which use mbr.
	copy_mbr($srci->{device}, $dsti->{device});
	apply_disk_info($dsti);

	if ($opts{partonly}) {
		print "no partition copy, due to --partonly\n";
		print "Succeeded.\n";
		exit 0;
	}

	copy_all_partitions_gpt($srci, $dsti) if (is_gpt($srci));
	copy_all_partitions_dos($srci, $dsti) if (is_dos($srci));

	print "Succeeded.\n";
	exit 0
}