#!/usr/bin/env ruby

# truncate -s +10G

# Usage: ruby extend_mbr.rb disk.img

SECTOR_SIZE = 512
MBR_SIZE = 512
PARTITION_ENTRY_OFFSET = 446
PARTITION_ENTRY_SIZE = 16
MAX_PARTITIONS = 4

if ARGV.length != 1
  puts "Usage: #{$0} disk.img"
  exit 1
end

image_path = ARGV[0]

# Read disk image and MBR
image_size = File.size(image_path)
total_sectors = image_size / SECTOR_SIZE

puts ">> Image size: #{image_size} bytes"
puts ">> Total sectors: #{total_sectors}"

File.open(image_path, 'rb+') do |f|
  mbr = f.read(MBR_SIZE).dup

  last_partition_index = nil
  start_sector = nil

  # Find last valid partition
  MAX_PARTITIONS.times do |i|
    offset = PARTITION_ENTRY_OFFSET + i * PARTITION_ENTRY_SIZE
    part_type = mbr.getbyte(offset + 4)

    if part_type != 0x00
      last_partition_index = i
    end
  end

  if last_partition_index.nil?
    abort "!! No valid partitions found in MBR"
  end

  entry_offset = PARTITION_ENTRY_OFFSET + last_partition_index * PARTITION_ENTRY_SIZE

  # Read start sector (little-endian uint32 at offset +8)
  start_sector_bytes = mbr.byteslice(entry_offset + 8, 4)
  start_sector = start_sector_bytes.unpack("V").first

  new_size = total_sectors - start_sector

  puts ">> Found partition ##{last_partition_index + 1}"
  puts "   Start sector: #{start_sector}"
  puts "   New size:     #{new_size} sectors"

  # Pack new size (little-endian uint32)
  new_size_bytes = [new_size].pack("V")
  mbr[entry_offset + 12, 4] = new_size_bytes

  # Write back modified MBR
  f.seek(0)
  f.write(mbr)

  puts "✅ MBR updated. Partition now fills remaining disk space."
end


# sudo e2fsck -f /dev/sda1
# sudo resize2fs /dev/sda1