#!/usr/bin/ruby

# Take a firmware file for various NAS devices, and write out it's
# constituent parts.
#
# Copyright (C) 2008 Matt Palmer <mpalmer@hezmatt.org>
#
#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation; either version 2 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program; if not, write to the Free Software
#   Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301 USA
#
# Script based on information provided by Leschinsky Oleg.
#
# Takes a firmware file image as it's sole mandatory argument, and verifies
# that the signature and checksums are OK before extracting whichever
# element(s) of the firmware you requested (if any).  Also displays the
# Product/Custom/Model ID, and the firmware signature type.

require 'optparse'

def main
	opts = parse_cmdline

	validate_command_line(opts)
	
	header = {}
	
	header[:k_offset],
	header[:k_size],
	header[:i_offset],
	header[:i_size],
	header[:d_offset],
	header[:d_size],
	header[:k_sum],
	header[:i_sum],
	header[:d_sum],
	header[:signature],
	header[:prod_id],
	header[:custom_id],
	header[:model_id],
	header[:unused_1],
	header[:unused_2],
	header[:unused_3],
	header[:unused_4] = File.read(opts[:image], 64).unpack("V9a12c5a7V")
	
	if header[:signature] == "\x55\xAAFrodoII\x00\x55\xAA"
		puts "FrodoII firmware signature found"
	elsif header[:signature] == "\x55\xAAChopper\x00\x55\xAA"
		puts "Chopper firmware signature found"
	elsif header[:signature] == "\x55\xAAGandolf\x00\x55\xAA"
		puts "Gandolf firmware signature found"
	else
		puts "ERROR: Unknown firmware signature (I got #{header[:signature].inspect}; perhaps it's new?)"
		exit 1
	end

	verify_sum(File.read(opts[:image], header[:k_size], header[:k_offset]), header[:k_sum], "Kernel")
	puts "Kernel is #{header[:k_size]} bytes"
	verify_sum(File.read(opts[:image], header[:i_size], header[:i_offset]), header[:i_sum], "initrd")
	puts "initrd is #{header[:i_size]} bytes"
	verify_sum(File.read(opts[:image], header[:d_size], header[:d_offset]), header[:d_sum], "defaults.tar.gz")
	puts "defaults.tar.gz is #{header[:d_size]} bytes"
	
	puts "Product ID: #{header[:prod_id]}"
	puts "Custom ID: #{header[:custom_id]}"
	puts "Model ID: #{header[:model_id]}"
	
	if opts[:kernel]
		File.open(opts[:kernel], 'w') do |fd|
			fd.write(File.read(opts[:image], header[:k_size], header[:k_offset]))
		end
		puts "Kernel data written to #{opts[:kernel]}"
		unless is_uboot(opts[:kernel])
			puts "WARNING: kernel data does not appear to be a uBoot file"
		end
	end
	
	if opts[:initrd]
		File.open(opts[:initrd], 'w') do |fd|
			fd.write(File.read(opts[:image], header[:i_size], header[:i_offset]))
		end
		puts "initrd data written to #{opts[:initrd]}"
		unless is_uboot(opts[:initrd])
			puts "WARNING: initrd data does not appear to be a uBoot file"
		end
	end
	
	if opts[:defaults]
		File.open(opts[:defaults], 'w') do |fd|
			fd.write(File.read(opts[:image], header[:d_size], header[:d_offset]))
		end
		puts "defaults.tar.gz written to #{opts[:defaults]}"
	end
end

def validate_command_line(opts)
	if !File.exist?(opts[:image])
		$stderr.puts "Firmware image file (#{opts[:image]}) doesn't exist!"
		exit 1
	end
end

def is_uboot(file)
	File.read(file, 4) == "\x27\x05\x19\x56"
end

def checksum(s)
	sum = 0
	

	sum
end

def verify_sum(s, expected_sum, desc)
	actual_sum = 0
	s.unpack("V*").each { |v| actual_sum ^= v }

	if actual_sum == expected_sum
		puts "#{desc} checksum matches"
	else
		puts "WARNING: #{desc} checksum mismatch (expected #{expected_sum}, got #{actual_sum})"
	end
end

def parse_cmdline
	opts = OptionParser.new
	optargs = {}

	opts.on('-h', '--help', "Print this help") { optargs[:help] = true }
	opts.on('-k KERNEL', '--kernel KERNEL', "Write out the kernel to the specified file", String) { |optargs[:kernel]| }
	opts.on('-i INITRD', '--initrd INITRD', "Write out the initrd to the specified file", String) { |optargs[:initrd]| }
	opts.on('-d DEFAULTS', '--defaults DEFAULTS', "Write out the defaults.tar.gz to the specified file", String) { |optargs[:defaults]| }

	image = opts.parse(ARGV)

	if optargs[:help]
		$stderr.puts opts.to_s
		exit 0
	end
	
	if image.nil? or image.empty?
		$stderr.puts "No firmware image provided!"
		exit 1
	end
	
	if image.length > 1
		$stderr.puts "I can only read one firmware image!"
		exit 1
	end
		
	optargs[:image] = image[0]

	optargs
end

main