#!/usr/bin/env perl
use strict;
use warnings;
use autodie qw(:all);
use Digest::CRC qw(crc);

die "Usage: $0 in_file [out_file_base]\n" if scalar @ARGV != 2 && scalar @ARGV != 1;
my $inp_file_name = $ARGV[0];
my $out_file_base = $ARGV[1];
open my $in, '<:raw', $inp_file_name;
seek $in, 0, 2;
my $file_size = tell $in;
seek $in, 0, 0;
read $in, my $header, 4160;

my ($magic, $fs, $version, $padding, $block_size, $device_size, $total_blocks, $used_blocks) = 
  unpack 'a15a15a4slQQQ', $header;
die "Wrong file magic\n" if $magic ne 'partclone-image';
die "Unsupported version\n" if $version ne '0001';

# The bitmap is 1 byte per block.
# Fixme: process the bitmap in chunks to lower memory usage.
read $in, (my $bitmap), $total_blocks;
my $used_blocks_from_bitmap = $bitmap =~ tr/\001/\001/;
die "Unexpected inconsistency: bitmap and header do not match ($used_blocks_from_bitmap != $used_blocks)\n"
  if $used_blocks_from_bitmap != $used_blocks;
read $in, $magic, 8;
die "Wrong magic marker after bitmap\n" if $magic ne 'BiTmAgIc';
my $expected_size = tell($in) + $used_blocks * ($block_size + 4);
if ($expected_size != $file_size) {
  print "Wrong image file size: expected $expected_size, got $file_size\n" 
}

{
  my $out_fh;
  my $i = 0;
  sub open_out {
    return if !defined $out_file_base;
    my $out_file_name = "$out_file_base.$i";
    print "Writing to file $out_file_name\n";
    open $out_fh, '>:raw', $out_file_name;
    $i++;
  }
  sub close_out {
    return if !defined $out_fh;
    close $out_fh;
  }
  sub out {
    return if !defined $out_fh;
    print $out_fh shift or die("Write error: $!\n");
  }
}

open_out;
out $header;
out $bitmap;
out $magic;

my $crc = 0xffffffff;

# Buggy partclone implementation of CRC
sub pc_crc {
  my ($crc, $data) = @_;

  # CRC32 is calculated from the first byte only due to a bug in the code.
  return crc(scalar (unpack 'a', $data) x $block_size, 32, $crc, 0, 1, 0x04C11DB7, 1, 1);
}

{
  my $forward_blocks = 4;

  sub forward_check {
    my $save_pos = tell $in;
    my $off = -1;

    TRY: {
      $off++;
      seek $in, $save_pos + $off - 4, 0;
      read $in, my $crc, 4;
      $crc = unpack 'L', $crc;
      for(1..$forward_blocks) {
        return undef if tell($in) + $block_size + 4 > $file_size;
        read $in, (my $data), $block_size;
        read $in, my $crc2, 4;
        $crc2 = unpack 'L', $crc2;
        $crc = pc_crc($crc, $data);
        goto TRY if $crc != $crc2;
      }
    }
    seek $in, $save_pos, 0;
    return $off;
  }
}

my @good;
push @good, 0;

print "Checking blocks...\n";

my $off;
while(1) {
  $off = tell $in;
  last if $off + $block_size + 4 > $file_size;
  read $in, (my $data), $block_size;
  read $in, my $crc2, 4;
  $crc2 = unpack 'L', $crc2;

  $crc = pc_crc($crc, $data);
  if ($crc != $crc2) {
    push @good, $off;
    print "Bad CRC at file offset $off, expected $crc, found $crc2\n";
    print "Searching forward for good blocks\n";
    seek $in, $off, 0;
    my $corr_bytes = forward_check();
    last if !defined $corr_bytes;
    $off += $corr_bytes;
    push @good, $off;
    print "Good blocks at offset $off\n";
    seek $in, $off, 0;
    read $in, $data, $block_size;
    read $in, $crc2, 4;
    $crc2 = unpack 'L', $crc2;
    $crc = $crc2;
    open_out;
  }
  out $data;
  out pack('L', $crc2);
  $used_blocks--;
}
push @good, $off;
my $missing_bytes = $used_blocks*($block_size+4);
print "Missing blocks: $used_blocks ($missing_bytes bytes)\n";
close_out;
print "Summary of good areas:\n";
while(scalar @good) {
  my $start = shift @good;
  my $end = shift @good;
  my $length = $end-$start;
  print "at $start length $length\n";
}