#!/usr/bin/perl
use strict;
use warnings;
use autodie qw(:all);
use Fcntl qw(SEEK_SET);

my ($badblock_file, $image_file, $part_offset) = @ARGV;

open my $image_file_h, '+<', $image_file;
sysseek $image_file_h, $part_offset, SEEK_SET;
my $boot_sect;
sysread $image_file_h, $boot_sect, 512;
die "FAT32 signature not found\n" if substr($boot_sect, 0x052, 8) ne 'FAT32   ';

my $sector_size = unpack 'v', substr($boot_sect, 11, 2);
my $reserved_sectors = unpack 'v', substr($boot_sect, 14, 2);
my $fats = unpack 'C', substr($boot_sect, 16, 1);
die "Only 1 FAT\n" if $fats < 2;
my $sectors_per_fat = unpack 'V', substr($boot_sect, 0x024, 4);
my $flags = unpack 'v', substr($boot_sect, 0x028, 2);
die "FATs not mirrored\n" if $flags & 1 << 6;

open my $badblock_file_h, '<', $badblock_file;
my $badblock_contents;
sysread $badblock_file_h, $badblock_contents, $sector_size;

my $fat_offset = $part_offset + $reserved_sectors * $sector_size;
my $fat_size = $sectors_per_fat * $sector_size;

sub get_fat_sector {
  my ($f, $i) = @_;
  sysseek $image_file_h, $fat_offset + $fat_size * $f + $sector_size * $i, SEEK_SET;
  my $sector_data;
  sysread $image_file_h, $sector_data, $sector_size;
  return $sector_data;
}

sub set_fat_sector {
  my ($f, $i, $data) = @_;
  sysseek $image_file_h, $fat_offset + $fat_size * $f + $sector_size * $i, SEEK_SET;
  syswrite $image_file_h, $data;
}

for(my $i = 0; $i < $sectors_per_fat; $i++) {
  my (@good, @bad);
  for (my $f = 0; $f < $fats; $f++) {
    my $sector_data = get_fat_sector($f, $i);
    if ($sector_data eq $badblock_contents) {
      push @bad, $f;
    } else {
      push @good, $f;
    }
  }
  if (scalar @bad > 0) {
    if (scalar @good > 0) {
      my $good_f = shift @good;
      my $data = get_fat_sector($good_f, $i);
      while(scalar @bad > 0) {
        my $bad_f = shift @bad;
        print "Repairing bad sector $i of FAT $bad_f with data from FAT $good_f\n";
        set_fat_sector($bad_f, $i, $data);
      }
    } else {
      print "Can't repair sector $i, all FATs bad\n";
    }
  }
}