#!/usr/bin/perl -CSDA

# vigil.pl - an improved md5 checksumming program with progress bar and colour.
# Copyright 2014 James Clark. All rights reserved.
# This program is Free software; you can redistribute it and/or modify it under the same terms as Perl itself.

# $Id: vigil.pl 626 2014-08-01 10:16:44Z james $

use strict;
use warnings;
use utf8;
use Encode qw/encode_utf8 decode_utf8/;
use Getopt::Long;
use Digest::MD5;
use Term::ANSIColor;
use Term::ProgressBar;
use File::Find;
use File::Basename;
use Data::Dumper;

# Define our command-line arguments.
my %opts = ( 'blocksize' => 16384 );
GetOptions(\%opts, "verify=s", "create=s", "update=s", "files", "blocksize=s", "help!");

# Fix up --blocksize argument if any
$opts{blocksize} *= 1024 if $opts{blocksize} =~ /^\d+[km]$/i;
$opts{blocksize} *= 1024 if $opts{blocksize} =~ /^\d+[m]$/i;
die "--blocksize must a number be greater than 0, optionally with k or m suffixes!" unless $opts{blocksize} > 0;


sub help
{
	print "Usage:-\n";
	print "  $0 --verify filename.md5sum\n";
	print "  $0 --create filename.md5sum <files and dirs to check>\n";
	print "  $0 --update filename.md5sum <files and dirs to check>\n";
	print "  $0 --files <files and dirs to check>\n";
	print "\n";
	print "     --blocksize <size>{k,m}\n";
	exit 1;
}


# plan entries are hashrefs with keys: filename, size (bytes), correct_md5 (optional),
# at verification end, additional keys are: computed_md5, start_time (unix timestamp), elapsed_time (secs).
# This way we can attempt an accurate time estimate.
sub create_plan_for_filename
{
	my ($filename, $correct_md5) = @_;
	
	my $plan_entry = {};
	$plan_entry->{filename} = $filename;
	
	if (-f $filename && -r $filename) {
		$plan_entry->{size} = -s $filename;
	} else {
		# Doesn't exist or is a directory or a socket or something. Die?
		$plan_entry->{error} = "Is not readable" unless -r $filename;
		$plan_entry->{error} = "Is not a plain file" unless -f $filename;
		$plan_entry->{error} = "Is a directory" if -d $filename;
		$plan_entry->{error} = "Does not exist" unless -e $filename;
	}
	
	# If the md5 is supplied by the caller, then we're checking a file. If not, we're probably creating one.
	if (defined $correct_md5) {
		$plan_entry->{correct_md5} = $correct_md5;
	}
	
	return $plan_entry;
}


# No matter whether we're verifying or creating a checksum, we're basically gonna run the same thing over all the
# files given to us in the plan: stream bytes through md5, report progress as we do so.
# Only once the plan has been executed do we *really* need to report which files don't match, although a running
# alert for failing files would be nice too.
sub run_md5_file
{
	my ($plan_entry, $progress_fn) = @_;
	
	# We use the OO interface to Digest::MD5 so we can feed it data a chunk at a time.
	my $md5 = Digest::MD5->new();
	my $current_bytes_read = 0;
	my $buffer;
	$plan_entry->{start_time} = time();
	$plan_entry->{elapsed_time} = 0;
	$plan_entry->{elapsed_bytes} = 0;
	
	# 3 argument form of open() allows us to specify 'raw' directly instead of using binmode and is a bit more modern.
	open(my $fh, '<:raw', $plan_entry->{filename}) or die "Couldn't open file $plan_entry->{filename}, $!\n";
	
	# Read the file in chunks and feed into md5.
	while ($current_bytes_read = read($fh, $buffer, $opts{blocksize})) {
		$md5->add($buffer);
		$plan_entry->{elapsed_bytes} += $current_bytes_read;
		$plan_entry->{elapsed_time} = time() - $plan_entry->{start_time};
		&$progress_fn($plan_entry->{elapsed_bytes});
	}
	# The loop will exit as soon as read() returns 0 or undef. 0 is normal EOF, undef indicates an error.
	die "Error while reading $plan_entry->{filename}, $!\n" if ( ! defined $current_bytes_read);
	
	close($fh) or die "Couldn't close file $plan_entry->{filename}, $!\n";
	
	# We made it out of the file alive. Store the md5 we computed. Note that this resets the Digest::MD5 object.
	$plan_entry->{computed_md5} = $md5->hexdigest();
}


# While run_md5_file is responsible for processing an individual file, run_md5_files_with_progress is the one
# that is responsible for doing a large batch of files, as well as giving visual feedback via a progress bar.
sub run_md5_files_with_progress
{
	my ($progressbar_title, @plan) = @_;

	# How big is this going be?
	my $total_size = get_plan_size(@plan);
	my $total_progress_so_far = 0;		# updated after each file is processed.
	
	# Initialise the progress bar.
	my $bar = Term::ProgressBar->new({
				name => $progressbar_title,
				count => $total_size,
				ETA => 'linear',
			});
	$bar->minor(0);		# No asterisks used for progress animation.
	my $next_bar_update = 0;		# used to ensure we don't spend more time drawing the bar than actual work.
	
	# Do the sums.
	foreach my $plan_entry (@plan) {
		next unless defined $plan_entry->{filename};
		next unless defined $plan_entry->{size};
		$bar->message($plan_entry->{filename});
		run_md5_file($plan_entry,
				sub {
					my $progress = $total_progress_so_far + $_[0];
					if ($progress > $next_bar_update) {
						$next_bar_update = $bar->update($progress);
					}
				});
		
		$total_progress_so_far += $plan_entry->{size};
		$next_bar_update = $bar->update($total_progress_so_far);
	}
	$bar->update($total_progress_so_far);
	print STDERR "\n";
	
	# Show the user.
	print_plan_sums(@plan);
	print_plan_pass_fail_summary(@plan);
	print_plan_errors(@plan);
}


sub verify_file_sums
{
	my ($md5sum_filename) = @_;
	
	# Initialise our plan from the md5sum file.
	my @plan = load_md5sum_file($md5sum_filename);
	
	# Go sum all those files.
	run_md5_files_with_progress("Verifying MD5", @plan);
}


sub create_file_sums
{
	my ($md5sum_filename, @dirs_and_files_to_sum) = @_;
	
	die "The file '$md5sum_filename' exists already and does not appear to be writable!\n" if (-e $md5sum_filename && ! -w $md5sum_filename);
	die "The dir '$md5sum_filename' is to be made in does not appear to be writable!\n" if ( ! -e $md5sum_filename && ! -w dirname($md5sum_filename));
	
	# Initialise our plan from the files and directories given on the commandline.
	my @plan = build_plan_from_directories(@dirs_and_files_to_sum);
	
	# Go sum all those files.
	run_md5_files_with_progress("Creating MD5", @plan);
	
	# The checksums we are making are all new, and we can only assume they are all correct.
	foreach my $plan_entry (@plan) {
		# If there isn't a {correct_md5} attribute, we should set it from the {computed_md5}.
		$plan_entry->{correct_md5} //= $plan_entry->{computed_md5};
	}
	
	# Then write the computed MD5s to the md5sum file given with --create.
	save_md5sum_file($md5sum_filename, @plan);
}


sub update_file_sums
{
	my ($md5sum_filename, @dirs_and_files_to_sum) = @_;
	
	die "The file '$md5sum_filename' does not exist!\n" if ( ! -e $md5sum_filename);
	die "The file '$md5sum_filename' does not appear to be writable!\n" if (-e $md5sum_filename && ! -w $md5sum_filename);

	# We need to load both the contents of our .md5sum file and build an index of the files in given directories,
	# so that we can compare them.
	my @official_plan = load_md5sum_file($md5sum_filename);
	my @actual_plan = build_plan_from_directories(@dirs_and_files_to_sum);
	
	my %official_filenames = map { $_->{filename}, $_ } grep { defined $_->{filename} } @official_plan;
	my %actual_filenames = map { $_->{filename}, $_ } grep { defined $_->{filename} } @actual_plan;
	
	my @new_files = grep { $_->{filename} && ! defined $official_filenames{$_->{filename}} } @actual_plan;
	my @removed_files = grep { $_->{filename} && ! defined $actual_filenames{$_->{filename}} } @official_plan;

	if (scalar @removed_files == scalar keys %official_filenames) {
		print STDERR colored("All files mentioned in '$md5sum_filename' would be removed - refusing to act, something doesn't look right here.", 'bold red'), "\n";
		print STDERR colored("If you really want to do this, use --create to create a brand-new .md5sum file instead.", 'bold red'), "\n";
		exit 1;
	}
	if (@removed_files == 0 && @new_files == 0) {
		print STDERR colored("No changes necessary.", 'bold green'), "\n";
		exit 0;
	}

	if (@removed_files) {
		print STDERR colored(scalar @removed_files . " Removed files:-", 'red'), "\n";
		print_plan_sums(@removed_files);
	}
	if (@new_files) {	
		print STDERR colored(scalar @new_files . " New files:-", 'cyan'), "\n";
	}
	# Go sum the new files.
	run_md5_files_with_progress("Updating MD5", @new_files);
	
	# Remove the removed files and add the new files. Leave unchanged entries as per the original.
	my @updated_plan;
	foreach my $plan_entry (@official_plan, @new_files) {
		# Skip files that can no longer be found.
		next unless $plan_entry->{filename} && defined $actual_filenames{$plan_entry->{filename}};
		
		# If there isn't a {correct_md5} attribute, we should set it from the {computed_md5}.
		$plan_entry->{correct_md5} //= $plan_entry->{computed_md5};
		
		push @updated_plan, $plan_entry;
	}

	print STDERR colored("Saving file '$md5sum_filename'...", 'bold green');
	save_md5sum_file($md5sum_filename, @updated_plan);
	print STDERR colored(" Done!", 'bold green'), "\n";
}


sub build_plan_from_directories
{
	my (@dirs_and_files_to_sum) = @_;
	# File::Find works with binary strings from the filesystem only.
	# This means our nice utf8 strings we get in @ARGV (via the -CA perl flag)
	# shouldn't be fed to find() - it must be concatenating them behind the
	# scenes and this causes Amusing Failure when attempting to traverse a directory
	# with utf8 code points. We must convert our UTF-X flagged wide character
	# Perl strings to utf-8 encoded binary data strings before feeding them to
	# the module, and utf8::decode() the filenames we get out of it as we traverse
	# the directory tree.
	my @dirs_and_files_to_sum_bin = map { encode_utf8($_); } @dirs_and_files_to_sum;
	
	# Initialise our plan from the files and directories given on the commandline.
	# Traversing a big tree recursively could take a while, so let the user know.
	print STDERR colored("Constructing plan...", 'yellow'), "\n";
	my @plan;
	find(	{
				# We pass in an anonymous sub to File::Find for it to run on every file and directory
				# as it traverses the tree. Note that the parameter name, 'wanted', is a misnomer;
				# File::Find will not do anything with the return value of our sub.
				wanted => sub {
					# Apparently File::Find doesn't handle utf8 filenames.
					my $fullpathname = decode_utf8($File::Find::name);
					
					if (-e $fullpathname && ! -d $fullpathname) {
						push @plan, create_plan_for_filename($fullpathname);
					}
				},
				follow_skip => 2,		# If we get duplicate dirs, files etc. during find, ignore and carry on.
				no_chdir => 1,			# We want our current dir consistent after Find does its work.
			},
			@dirs_and_files_to_sum_bin
		);
	@plan = prune_plan_duplicates(@plan);
	print STDERR colored("Complete. Plan contains " . scalar @plan . " items.", 'yellow'), "\n";
	
	return @plan;
}


sub display_file_sums
{
	my (@filenames) = @_;
	
	# Initialise our plan from the filenames given on the commandline.
	my @plan = map { create_plan_for_filename($_) } @filenames;
	
	# Go sum all those files.
	run_md5_files_with_progress("Computing MD5", @plan);
}


sub prune_plan_duplicates
{
	my (@plan) = @_;
	my @uniq_plan;
	my %names;
	foreach my $plan_entry (@plan) {
		next unless $plan_entry->{filename};
		if ($names{$plan_entry->{filename}} // 0 == 1) {
			print STDERR colored("Warning: File '$plan_entry->{filename}' mentioned multiple times. Duplicate entries in plan will be ignored.", 'bold yellow'), "\n";
		}
		push @uniq_plan, $plan_entry unless $names{$plan_entry->{filename}};
		$names{$plan_entry->{filename}}++;
	}
	return @uniq_plan;
}


sub print_plan_sums
{
	my (@plan) = @_;
	foreach my $plan_entry (@plan) {
		next unless $plan_entry->{filename};
		
		my $sumdata = $plan_entry->{computed_md5} // $plan_entry->{correct_md5} // "(none)";
		my $sumcolour = "reset";
		if ($plan_entry->{computed_md5}) {
			$sumcolour = "cyan";
		}
		if ($plan_entry->{computed_md5} && $plan_entry->{correct_md5}) {
			$sumcolour = ( $plan_entry->{computed_md5} eq $plan_entry->{correct_md5} ) ? "green" : "red";
		}
		
		# Only output colour if we're on a TTY; user might want to > to a file.
		if (-t STDOUT) {
			print colored($sumdata, $sumcolour) . "  " . $plan_entry->{filename}, "\n";
		} else {
			print "$sumdata  $plan_entry->{filename}\n";
		}
	}
}


sub print_plan_pass_fail_summary
{
	my (@plan) = @_;
	my @checked = grep { $_->{computed_md5} } @plan;
	my @verified = grep { $_->{correct_md5} && ( $_->{computed_md5} // "fail" ) eq $_->{correct_md5} } @plan;
	my @failures = grep { $_->{correct_md5} && ( $_->{computed_md5} // "fail" ) ne $_->{correct_md5} } @plan;
	my $full_size = get_plan_size(@plan);
	my $elapsed_bytes = get_plan_elapsed_bytes(@plan);
	my $elapsed_time = get_plan_elapsed_time(@plan);
	my $avg_speed = $elapsed_bytes / ($elapsed_time || 1);
	
	if (@plan > 0) {
		print STDERR colored(scalar @plan . " files considered, " . scalar @checked . " checked.", 'cyan'), "\n";
	}
	if (@plan > 0 && $elapsed_bytes > 0) {
		$full_size = format_byte_size($full_size);
		$elapsed_bytes = format_byte_size($elapsed_bytes);
		$elapsed_time = format_time_duration($elapsed_time);
		$avg_speed = format_byte_size($avg_speed);
		print STDERR colored("$elapsed_bytes of $full_size checked in $elapsed_time ($avg_speed/s)", 'cyan'), "\n";
	}
	if (@verified > 0) {
		print STDERR colored(scalar @verified . " files verified OK.", 'green'), "\n";
	}
	if (@failures > 0) {
		print STDERR colored(scalar @failures . " files failed checksum:-", 'bold red'), "\n";
		foreach my $failure (@failures) {
			print STDERR colored("   " . $failure->{filename}, 'reset'), "\n";
		}
	}
}


sub print_plan_errors
{
	my (@plan) = @_;
	my %files_by_error;
	foreach my $plan_entry (@plan) {
		next unless $plan_entry->{error};
		$files_by_error{$plan_entry->{error}} //= [];
		push @{ $files_by_error{$plan_entry->{error}} }, $plan_entry;
	}
	if (keys %files_by_error > 0) {
		print STDERR colored("Errors were encountered during processing:-", 'bold red'), "\n";
		foreach my $error (keys %files_by_error) {
			print STDERR colored("   [$error]:-", 'bold red'), "\n";
			foreach my $failure (@{$files_by_error{$error}}) {
				print STDERR colored("      " . $failure->{filename}, 'reset'), "\n";
			}
		}
	}
}


sub get_plan_size
{
	return get_plan_summed_field_value('size', @_);
}

sub get_plan_elapsed_bytes
{
	return get_plan_summed_field_value('elapsed_bytes', @_);
}

sub get_plan_elapsed_time
{
	return get_plan_summed_field_value('elapsed_time', @_);
}



sub get_plan_summed_field_value
{
	my ($fieldname, @plan) = @_;
	my $total = 0;
	foreach my $entry (@plan) {
		$total += $entry->{$fieldname} // 0;
	}
	return $total;
}


sub format_byte_size
{
	my ($bytes) = @_;
	my $unit = 'B';
	if ($bytes > 1024 * 1024 * 1024) {
		$bytes /= 1024 * 1024 * 1024;
		$unit = 'GiB';
	} elsif ($bytes > 1024 * 1024) {
		$bytes /= 1024 * 1024;
		$unit = 'MiB';
	} elsif ($bytes > 1024) {
		$bytes /= 1024;
		$unit = 'KiB';
	}
	return sprintf("%.01f %s", $bytes, $unit);
}


sub format_time_duration
{
	my ($total_secs) = @_;
	my $days = "";
	my $hours = "";
	my $mins = "";
	my $secs = "";
	if ($total_secs > 60 * 60 * 24) {
		$days = int($total_secs / (60 * 60 * 24)) . "d";
		$total_secs %= 60 * 60 * 24;
	}
	if ($total_secs > 60 * 60) {
		$hours = int($total_secs / (60 * 60)) . "h";
		$total_secs %= 60 * 60;
	}
	if ($total_secs > 60) {
		$mins = int($total_secs / 60) . "m";
		$total_secs %= 60;
	}
	$secs = $total_secs . "s";
	return "$days$hours$mins$secs";
}


sub load_md5sum_file
{
	my ($filename) = @_;
	my @plan;
	
	# Yep, we're assuming everything is utf8. Things which are not utf8 cause you pain. Horrible icky pain.
	open(my $fh, '<:utf8', $filename) or die "Couldn't open '$filename' : $!\n";
	my $linenum = 0;
	while (my $line = <$fh>) {
		chomp $line;
		$linenum++;
		if ($line =~ /^(?<md5>\p{ASCII_Hex_Digit}{32})  (?<filename>.*)$/) {
			# Checksum and filename compatible with md5sum output.
			push @plan, create_plan_for_filename($+{filename}, $+{md5});
			
		} elsif ($line =~ /^(?<md5>\p{ASCII_Hex_Digit}{32})  (?<filename>.*)$/) {
			# Checksum and filename compatible with md5sum's manpage but not valid for the actual program.
			# We'll use it, but complain.
			print STDERR colored("Warning: ", 'bold red'), colored("md5sum entry '", 'red'), $line, colored("' on line $linenum of file $filename is using only one space, not two - this doesn't match the output of the actual md5sum program!.", 'red'), "\n";
			push @plan, create_plan_for_filename($+{filename}, $+{md5});
			
		} elsif ($line =~ /^\s*$/) {
			# Blank line, ignore.
			
		} else {
			# No idea. Best not to keep quiet, it could be a malformed checksum line and we don't want to just quietly skip the file if so.
			print STDERR colored("Warning: ", 'bold red'), colored("Unrecognised md5sum entry '", 'red'), $line, colored("' on line $linenum of file $filename.", 'red'), "\n";
			push @plan, { error => "Unrecognised md5sum entry" };
		}
	}
	close($fh) or die "Couldn't close '$filename' : $!\n";
	
	return @plan;
}


sub save_md5sum_file
{
	my ($filename, @plan) = @_;
	
	my $fh;
	unless (open($fh, '>:utf8', $filename)) {
		print STDERR colored("Warning: Couldn't write to '$filename' : $!", 'bold red'), "\n";
		$filename = "/tmp/vigil.$$." . basename($filename);
		print STDERR colored("         Attempting to write to '$filename' instead.", 'red'), "\n";
		open($fh, '>:utf8', $filename) or die "Argh, that failed as well! : $!\n";
	}
	foreach my $plan_entry (@plan) {
		next unless $plan_entry->{correct_md5};
		next unless $plan_entry->{filename};
		print $fh "$plan_entry->{correct_md5}  $plan_entry->{filename}\n";
	}
	close($fh) or die "Couldn't close '$filename' : $!\n";
}


sub main
{
	if ($opts{files}) {
		# Just print out the sums of all the files given as @ARGV.
		display_file_sums(@ARGV);
		
	} elsif ($opts{create}) {
		# Compute the sums of all the files given as @ARGV or found recursively within directories given;
		# Create a new .md5sum file with the name specified after --create.
		create_file_sums($opts{create}, @ARGV);

	} elsif ($opts{update}) {
		# Compute the sums of all the files given as @ARGV or found recursively within directories given;
		# But only if they're files not already mentioned in the .md5sum file.
		# Also take note of any files that are mentioned in the file but are now absent.
		# Create a new .md5sum file with updated checksums for the new files and remove absent file sums.
		update_file_sums($opts{update}, @ARGV);

	} elsif ($opts{verify}) {
		# Compute the sums of all the files listed in the .md5sum file we get from --verify.
		# Tell us if there are mismatches.
		verify_file_sums($opts{verify});

	} else {
		help();
	}
}


help() if $opts{help};
main();

