#!/usr/bin/ruby

# Create a firmware file suitable for upload to a DNS-323 or CH3SNAS
# device.
#
# 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.
#
# To construct a firmware image, you need six things: a uBoot-packaged
# kernel, a uBoot-packaged initrd, the product/custom/model IDs for the
# particular device you want to target, and the firmware signature type.
# Using the wrong values can result in the device rejecting the firmware.
#
# Known values for these the IDs and firmware signature type are:
#
#                 prod   custom   model    fwtype
# DNS-323 rev B1:    7        1       1   FrodoII
# CH3SNAS:           7        2       1   FrodoII
# DNS-343:           9        1       1   Gandolf
# DNS-321:          10        1       1   Chopper
#

require 'optparse'

def main
	opts = parse_cmdline

	validate_command_line(opts)
	
	k_size = File.stat(opts[:kernel]).size
	i_size = File.stat(opts[:initrd]).size
	d_size = opts[:defaults].nil? ? 0 : File.stat(opts[:defaults]).size
	
	ctl_header = [
		64,
		k_size,
		64 + k_size,
		i_size,
		64 + k_size + i_size,
		d_size,
		checksum(opts[:kernel]),
		checksum(opts[:initrd]),
		d_size == 0 ? 0 : checksum(opts[:defaults]),
		"\x55\xAA#{opts[:signature]}\x00\x55\xAA",
		opts[:prod_id],
		opts[:custom_id],
		opts[:model_id],
		1,
		3,
		"\x00" * 7,
		0
	].pack("V9a12c5a7V")

	File.open(opts[:output], 'w') do |fd|
		fd.write ctl_header
		fd.write File.read(opts[:kernel])
		fd.write File.read(opts[:initrd])
		fd.write File.read(opts[:defaults]) unless d_size == 0
	end
	
	puts "Firmware generation completed successfully."
end

def validate_command_line(opts)
	%w{kernel initrd}.each do |k|
		if !File.exist?(opts[k.to_sym])
			$stderr.puts "#{k} file (#{opts[k.to_sym]}) doesn't exist!"
			exit 1
		end
	
		if !is_uboot(opts[k.to_sym])
			$stderr.puts "#{k} file #{opts[k.to_sym]} is not a uboot file"
			exit 1
		end
	end
	
	if opts[:defaults] and !File.exist?(opts[:defaults])
		$stderr.puts "default file (#{opts[:defaults]}) doesn't exist!"
		exit 1
	end
	
	if opts[:signature]
		unless %w{FrodoII Chopper Gandolf}.include? opts[:signature]
			$stderr.puts "Unknown signature type #{opts[:signature]}"
			exit 1
		end
	else
		opts[:signature] = "FrodoII"
	end
end

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

def checksum(file)
	sum = 0
	
	File.read(file).unpack("V*").each { |v| sum ^= v }

	sum
end

def parse_cmdline
	opts = OptionParser.new
	optargs = {}

	opts.on('-h', '--help', "Print this help") { optargs[:help] = true }
	opts.on('-k KERNEL', '--kernel KERNEL', "Specify the kernel to include in the firmware image", String) { |optargs[:kernel]| }
	opts.on('-i INITRD', '--initrd INITRD', "Specify the initrd to include in the firmware image", String) { |optargs[:initrd]| }
	opts.on('-d DEFAULTS', '--defaults DEFAULTS', "Specify the defaults.tar.gz to include in the firmware image (optional)", String) { |optargs[:defaults]| }
	opts.on('-o OUTPUT', '--output OUTPUT', "Specify where to put the resulting firmware image", String) { |optargs[:output]| }
	opts.on('-p PROD_ID', '--product-id PROD_ID', "The product ID to embed in the firmware image", Integer) { |optargs[:prod_id]| }
	opts.on('-c CUSTOM_ID', '--custom-id CUSTOM_ID', "The custom ID to embed in the firmware image", Integer) { |optargs[:custom_id]| }
	opts.on('-m MODEL_ID', '--model-id MODEL_ID', "The model ID to embed in the firmware image", Integer) { |optargs[:model_id]| }
	opts.on('-s SIGNATURE', '--signature SIGNATURE', "The firmware signature type (either FrodoII, Chopper or Gandolf)", String) { |optargs[:signature]| }

	opts.parse(ARGV)

	if optargs[:help]
		$stderr.puts opts.to_s
		exit 0
	end

	%w{kernel initrd output prod_id custom_id model_id}.each do |k|
		if optargs[k.to_sym].nil?
			$stderr.puts "Missing required argument #{k}"
			exit 1
		end
	end

	optargs
end

main