#!/usr/bin/ruby

# Create and dismantle firmware files suitable for upload to a range of
# D-Link (and compatible) NAS devices.
#
# Copyright (C) 2008,2012,2014 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
#
# This tool is based on information provided by Leschinsky Oleg.  Many
# thanks go to him for deciphering (and publishing) his analysis of the file
# formats involved.
#
# This file should not be called as dns323fw-tool; create links to this file
# named 'mkdns323fw' and 'splitdns323fw'.

require 'optparse'

begin
	require 'ffi'
rescue LoadError
	$stderr.puts "The 'ffi' gem is not available."
	$stderr.puts "Perhaps try running `gem install ffi`"
	exit 1
end

# This class handles all the grunt work for decoding and encoding firmware
# files.
class DnsFirmware
	# Represents a file which is a part of the firmware blob
	class IncludedFile
		# The contents of the file
		attr_reader :contents

		# The stored / calculated checksum
		attr_reader :checksum

		# This class can be instantiated in two different ways:
		#
		#  * By passing in `:filename => "some file name"` as the second
		#    argument; or
		#
		#  * By passing in `:contents => "some file contents"` and `:checksum
		#    => <number>` as the second argument.
		#
		def initialize(opts)
			if opts[:filename] and (opts[:contents] or opts[:checksum])
				raise ArgumentError,
				      "Must pass either :filename or :contents and :checksum to #{self.class}.initialize"
			end

			if opts[:filename]
				@filename = File.expand_path(opts[:filename])
				@contents = File.read(@filename).force_encoding("BINARY")
			elsif opts[:contents]
				@contents = opts[:contents]
				if opts[:checksum]
					@checksum = opts[:checksum]
				end
			else
				raise ArgumentError,
				      "Must pass exactly one of :filename or :contents to #{self.class}.initialize"
			end

			# Calculate a checksum if one was not provided
			@checksum ||= calculate_checksum
		end

		def validate
			if @filename and !File.exists?(@filename)
				log_error "does not exist"
			end

			unless @checksum == self.calculate_checksum
				$stderr.puts "Checksum of #{type} file is incorrect (expected #{@checksum}, got #{checksum(@contents)}"
			end

			true
		end

		def log_error(s)
			if @filename
				$stderr.print "#{@filename}: "
			end
			$stderr.puts s
		end

		def size
			@contents.length
		end

		def calculate_checksum
			@contents.unpack("V*").inject(0) { |v, i| v ^= i }
		end

		def write(f)
			File.write(f, @contents)
		end
	end

	# Any included file that should be a uBoot file
	class UBootFile < IncludedFile
		def validate
			super or return false

			unless @contents[0..3] == "\x27\x05\x19\x56"
				$stderr.puts "#{type} file #{@filename} does not appear to be a uBoot image file"
				return false
			end

			true
		end
	end

	class KernelFile < UBootFile
		def type
			:kernel
		end
	end

	class InitrdFile < UBootFile
		def type
			:initrd
		end
	end

	class DefaultsFile < IncludedFile
		def type
			:defaults
		end

		def validate
			super or return false

			unless @contents[0..1] == "\x1f\x8b"
				$stderr.puts "defaults file does not appear to be gzipped"
				return false
			end

			true
		end
	end

	class SquashfsFile < IncludedFile
		def type
			:squashfs
		end

		def validate
			super or return false

			unless @contents[0x800..0x803] == "hsqs" or
			       @contents[0x800..0x803] == "shsq"
				$stderr.puts "filesystem image file does not appear to be a squashfs filesystem"
				return false
			end

			true
		end
	end

	class GenericHeader < FFI::Struct
		def initialize(content = nil)
			super()

			if content
				if content.length != self.size
					raise ArgumentError,
					      "Content passed to #{self.class}.new must be #{self.size} bytes"
				end
				self.to_ptr.write_string(content)
			end
		end

		def data
			self.to_ptr.read_string(self.size)
		end
	end

	class FrodoHeader < GenericHeader
		layout :kernel_offset,   :uint32,      # 0
		       :kernel_size,     :uint32,      # 4
		       :initrd_offset,   :uint32,      # 8
		       :initrd_size,     :uint32,      # 12
		       :defaults_offset, :uint32,      # 16
		       :defaults_size,   :uint32,      # 20
		       :kernel_sum,      :uint32,      # 24
		       :initrd_sum,      :uint32,      # 28
		       :defaults_sum,    :uint32,      # 32
		       :magic1,          [:uint8, 2],  # 36
		       :sig,             [:uint8, 8],  # 38
		       :magic2,          [:uint8, 2],  # 46
		       :product_id,      :uint8,       # 48
		       :custom_id,       :uint8,       # 49
		       :model_id,        :uint8,       # 50
		       :compat_flag,     :uint8,       # 51
		       :compat_id,       :uint8,       # 52
		       :_1,              [:uint8, 11]  # 53

		def signature
			self[:sig].to_ptr.read_string.strip
		end

		def signature=(s)
			if s.length > 8
				raise ArgumentError,
				      "Signature can be no more than 8 bytes"
			end
			self[:sig].to_ptr.write_string(s)
		end

		def valid?
			self[:magic1].to_s == "\x55\xAA".force_encoding("BINARY") and
			self[:magic2].to_s == "\x55\xAA".force_encoding("BINARY")
		end
	end

	class LongHeader < GenericHeader
		layout :kernel_offset,   :uint32,      # 0
		       :kernel_size,     :uint32,      # 4
		       :initrd_offset,   :uint32,      # 8
		       :initrd_size,     :uint32,      # 12
		       :squashfs_offset, :uint32,      # 16
		       :squashfs_size,   :uint32,      # 20
		       :defaults_offset, :uint32,      # 24
		       :defaults_size,   :uint32,      # 28
		       :kernel_sum,      :uint32,      # 32
		       :initrd_sum,      :uint32,      # 36
		       :squashfs_sum,    :uint32,      # 40
		       :defaults_sum,    :uint32,      # 44
		       :magic1,          [:uint8, 2],  # 48
		       :sig,             [:uint8, 8],  # 50
		       :magic2,          [:uint8, 2],  # 58
		       :product_id,      :uint8,       # 60
		       :custom_id,       :uint8,       # 61
		       :model_id,        :uint8,       # 62
		       :compat_flag,     :uint8,       # 63
		       :compat_id,       :uint8,       # 64
		       :_1,              [:uint8, 63]  # 65

		def signature
			self[:sig].to_ptr.read_string.strip
		end

		def signature=(s)
			if s.length > 8
				raise ArgumentError,
				      "Signature can be no more than 8 bytes"
			end
			self[:sig].to_ptr.write_string(s)
		end

		def valid?
			self[:magic1].to_s == "\x55\xAA".force_encoding("BINARY") and
			self[:magic2].to_s == "\x55\xAA".force_encoding("BINARY")
		end
	end

	HEADER_FORMATS = [FrodoHeader, LongHeader]

	# This class can be initialized two ways:
	#
	# - with a single string argument, which is the filename of an existing
	#   firmware file to be dissected; or
	#
	# - with a hash containing the following keys:
	#   * `:kernel_file` (required) -- A file containing a raw bzImage kernel
	#
	#   * `:initrd_file` (required) -- A file containing an initrd
	#
	#   * `:defaults_file` (optional) -- A file containing a compressed
	#         tarball of configuration data; take a look at an existing
	#         firmware file for your device to see what can be in there.
	#
	#   * `:squashfs_file` (optional) -- A file containing a squashfs
	#         filesystem image.  Only some newer devices (DNS-320, etc)
	#         make use of such a file.
	#
	#   * `:product_id` (required) -- The "product ID" to embed in the
	#         firmware file.  What this needs to be set to is entirely
	#         dependent on your device and what it expects.  It *appears*
	#         that hardware that is substantially identical will tend to
	#         share a common product ID, even if it is badged differently.
	#
	#   * `:custom_id` (required) -- The "customisation ID" to embed in the
	#         firmware file.  What this needs to be set to is entirely
	#         dependent on your device and what it expects.  It appears that
	#         hardware that is substantially identical, but shipped by
	#         different vendors, will usually have a different custom ID.
	#
	#   * `:model_id` (required) -- The "model ID" to embed in the firmware
	#         file.  What this needs to be set to is entirely dependent on
	#         your device and what it expects.
	#
	#   * `:compat_id` (optional) -- The "compatibility version ID" to embed
	#         in the firmware file.  As far as can be determined, some
	#         firmware loaders will not load a firmware file which contains a
	#         compatibility ID lower than an arbitrary figure.
	#
	#   * `:signature` (required) -- The signature to embed in the firmware
	#         file.  This is some sort of magic string that lets the firmware
	#         loader know that its reading a real firmware file, as opposed
	#         to some random gibberish.  What you need to set this to is
	#         determined entirely by what your device expects to see.  The
	#         signature is also used to determine what kind of header to place
	#         in front of the firmware.
	#
	def initialize(opts)
		@files  = {}

		if opts.is_a? String
			@firmware_file = opts

			candidate_formats = HEADER_FORMATS.select do |hf|
				h = hf.new(File.read(@firmware_file, hf.size))
				h.valid?
			end

			if candidate_formats.empty?
				$stderr.puts "Unable to recognise the format of firmware file #{@firmware_file}"
				$stderr.puts "The file may be corrupted, or this firmware may be a format I am not familiar with."
				exit 1
			end

			if candidate_formats.length > 1
				$stderr.puts "This firmware is of an ambiguous format."
				$stderr.puts "Please inform theshed+dns323-firmware-tools@hezmatt.org, and include a"
				$stderr.puts "link to the firmware file for testing."
				exit 1
			end

			@header = candidate_formats.first.new(File.read(@firmware_file, candidate_formats.first.size))

			@files[:kernel] = KernelFile.new(
			                               :contents => File.read(
			                                                   @firmware_file,
			                                                   @header[:kernel_size],
			                                                   @header[:kernel_offset]
			                                                 ),
			                               :checksum => @header[:kernel_sum]
			                             )
			@files[:initrd] = InitrdFile.new(
			                               :contents => File.read(
			                                                   @firmware_file,
			                                                   @header[:initrd_size],
			                                                   @header[:initrd_offset]
			                                                 ),
			                               :checksum => @header[:initrd_sum]
		                             )
			if @header[:defaults_size] > 0
				@files[:defaults] = DefaultsFile.new(
				                                   :contents => File.read(
				                                                       @firmware_file,
				                                                       @header[:defaults_size],
				                                                       @header[:defaults_offset]
				                                                     ),
				                                   :checksum => @header[:defaults_sum]
				                                 )
			end
			if @header.members.include?(:squashfs_size) and @header[:squashfs_size] > 0
				@files[:squashfs] = SquashfsFile.new(
				                                   :contents => File.read(
				                                                       @firmware_file,
				                                                       @header[:squashfs_size],
				                                                       @header[:squashfs_offset]
				                                                     ),
				                                   :checksum => @header[:squashfs_sum]
				                                 )
			end
		elsif opts.is_a? Hash
			[:kernel_file, :initrd_file].each do |f|
				unless opts[f]
					raise ArgumentError,
					      "No #{f.inspect} provided"
				end
			end

			@files[:kernel] = KernelFile.new(:filename => opts[:kernel_file])
			@files[:initrd] = InitrdFile.new(:filename => opts[:initrd_file])
			if opts[:defaults_file]
				@files[:defaults] = DefaultsFile.new(:filename => opts[:defaults_file])
			end
			if opts[:squashfs_file]
				@files[:squashfs] = SquashfsFile.new(:filename => opts[:squashfs_file])
			end

			@product_id    = opts[:product_id]    or raise ArgumentError.new("No :product_id provided")
			@custom_id     = opts[:custom_id]     or raise ArgumentError.new("No :custom_id provided")
			@model_id      = opts[:model_id]      or raise ArgumentError.new("No :model_id provided")
			@compat_id     = opts[:compat_id]
			@signature     = opts[:signature]     or raise ArgumentError.new("No :signature provided")

			# This logic will probably have to be improved somewhat, likely via
			# another option, in the future
			if @signature =~ /^DNS32/
				@header = LongHeader.new
			else
				@header = FrodoHeader.new
			end

			@header[:magic1] = @header[:magic2] = "\x55\xAA"
		else
			raise ArgumentError.new("Incorrect type passed to DnsFirmware#initialize.  String or Hash expected, got #{opts.class}")
		end
	end

	def product_id
		@header[:product_id]
	end

	def custom_id
		@header[:custom_id]
	end

	def model_id
		@header[:model_id]
	end

	def compat_id
		@header[:compat_flag] ? @header[:compat_id] : nil
	end

	def kernel
		@files[:kernel]
	end

	def initrd
		@files[:initrd]
	end

	def defaults
		@files[:defaults]
	end

	def squashfs
		@files[:squashfs]
	end

	# Return the signature of this firmware file.
	def signature
		@header.signature
	end

	# This method works from the kernel/initrd/defaults/etc data and writes out
	# a complete firmware file to the destfile of your choosing.
	def write_firmware_file(destfile)
		@header.signature         = @signature
		@header[:product_id]      = @product_id
		@header[:custom_id]       = @custom_id
		@header[:model_id]        = @model_id
		@header[:compat_flag]     = @compat_id.nil? ? 0 : 1
		@header[:compat_id]       = @compat_id.to_i

		file_offset = @header.size
		file_types = @header.
		               members.
		               grep(/_sum$/).
		               map { |m| m.to_s.sub(/_sum$/, '').to_sym }


		file_types.each do |f_type|
			@header["#{f_type}_offset".to_sym] = file_offset
			if @files[f_type]
				@header["#{f_type}_size".to_sym] = @files[f_type].size
				@header["#{f_type}_sum".to_sym] = @files[f_type].checksum
				file_offset += @files[f_type].size
			else
				@header["#{f_type}_size".to_sym] = 0
				@header["#{f_type}_sum".to_sym] = 0
			end
		end

		File.open(destfile, 'w') do |fd|
			fd.write @header.data

			file_types.each do |f_type|
				fd.write @files[f_type].contents if @files[f_type]
			end
		end
	end

	def validate
		@files.values.each { |f| f.validate }
	end
end

def mkdns323fw(args)
	devices = {
		'DNS-325' => {
			:signature => 'DNS-325',
			:product_id   => 0,
			:custom_id => 8,
			:model_id  => 5
		},
		'DNS-323' => {
			:signature  => 'FrodoII',
			:product_id => 7,
			:custom_id  => 1,
			:model_id   => 1
		},
		'DNS-321' => {
			:signature  => 'Chopper',
			:product_id => 10,
			:custom_id  => 1,
			:model_id   => 1
		},
		'DNS-320' => {
			:signature  => 'DNS323D1',  # No, seriously...
			:product_id => 0,
			:custom_id  => 8,
			:model_id   => 7
		},
		'DNS-320B' => {
			:signature  => 'DNS320B',
			:product_id => 0,
			:custom_id  => 8,
			:model_id   => 12
		},
		'DNS-320L' => {
			:signature  => 'DNS320L',
			:product_id => 0,
			:custom_id  => 8,
			:model_id   => 11
		},
		'CH3SNAS' => {
			:signature  => 'FrodoII',
			:product_id => 7,
			:custom_id  => 2,
			:model_id   => 1
		},
	}

	opts = OptionParser.new
	optargs = {}

	opts.on('-h', '--help',
	        "Print this help") { puts opts.to_s; exit 0 }
	opts.on('-k KERNEL', '--kernel KERNEL',
	        "Specify the kernel to include in the firmware image",
	        String) { |k| optargs[:kernel_file] = k }
	opts.on('-i INITRD', '--initrd INITRD',
	        "Specify the initrd to include in the firmware image",
	        String) { |i| optargs[:initrd_file] = i }
	opts.on('-d DEFAULTS', '--defaults DEFAULTS',
	        "Specify the defaults.tar.gz to include in the firmware image (optional)",
	        String) { |d| optargs[:defaults_file] = d }
	opts.on('--squashfs SQUASHFS',
	        "Specify a squashfs filesystem image to include in the firmware image (optional)",
	        String) { |d| optargs[:squashfs_file] = d }
	opts.on('-o OUTPUT', '--output OUTPUT',
	        "Specify where to put the resulting firmware image",
	        String) { |o| optargs[:output] = o }
	opts.on('-p PROD_ID', '--product-id PROD_ID',
	        "The product ID to embed in the firmware image",
	        Integer) { |p| optargs[:product_id] = p }
	opts.on('-c CUSTOM_ID', '--custom-id CUSTOM_ID',
	        "The custom ID to embed in the firmware image",
	        Integer) { |c| optargs[:custom_id] = c }
	opts.on('-m MODEL_ID', '--model-id MODEL_ID',
	        "The model ID to embed in the firmware image",
	        Integer) { |m| optargs[:model_id] = m }
	opts.on('-s SIGNATURE', '--signature SIGNATURE',
	        "The firmware signature type",
	        String) { |s| optargs[:signature] = s }
	opts.on('-t DEVICE', '--device-type DEVICE',
	        "Specify all the IDs and signature by choosing a device",
	        "Known devices: #{devices.keys.join(' ')}",
	        String) { |s| optargs[:device] = s }

	opts.parse(args)

	if optargs[:device]
		optargs.merge!(devices[optargs.delete(:device)])
	end

	opts = optargs

	%w{kernel_file initrd_file output product_id custom_id model_id}.each do |k|
		if opts[k.to_sym].nil?
			$stderr.puts "Missing required argument #{k}"
			exit 1
		end
	end

	opts[:signature] ||= "FrodoII"

	# We just hardwire this to be the maximum value for now, since there is
	# no indication that it does any harm in the wild
	opts[:compat_id] = 255

	begin
		fw = DnsFirmware.new(opts)
		fw.validate and fw.write_firmware_file(opts[:output])
	rescue StandardError => e
		$stderr.puts "Firmware generation failed: #{e.class}: #{e.message}"
		e.backtrace.each { |l| puts "   #{l}" }
	else
		puts "Firmware generation completed successfully."
	end
end

def splitdns323fw(args)
	opts = OptionParser.new
	optargs = {}

	opts.on('-h', '--help',
	        "Print this help") { $stderr.puts opts.to_s; exit 0 }
	opts.on('-k KERNEL', '--kernel KERNEL',
	        "Write out the kernel to the specified file",
	        String) { |k| optargs[:kernel] = k }
	opts.on('-i INITRD', '--initrd INITRD',
	        "Write out the initrd to the specified file",
	        String) { |i| optargs[:initrd] = i }
	opts.on('-d DEFAULTS', '--defaults DEFAULTS',
	        "Write out the defaults.tar.gz to the specified file",
	        String) { |d| optargs[:defaults] = d }
	opts.on('-s SQUASHFS', '--squashfs SQUASHFS',
	        "Write out the squashfs filesystem image to the specified file",
	        String) { |s| optargs[:squashfs] = s }

	image = opts.parse(args)

	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

	image = image[0]

	fw = DnsFirmware.new(image)

	puts "'#{fw.signature}' firmware signature found"

	fw.kernel.validate or $stderr.puts "Kernel data failed validation"
	puts "Kernel is #{fw.kernel.size} bytes"
	fw.initrd.validate or $stderr.puts "Initrd data failed validation"
	puts "initrd is #{fw.initrd.size} bytes"
	if fw.defaults
		fw.defaults.validate or $stderr.puts "Defaults data failed validation"
		puts "defaults.tar.gz is #{fw.defaults.size} bytes"
	else
		puts "No defaults.tar.gz in this firmware file"
	end

	if fw.squashfs
		fw.squashfs.validate or $stderr.puts "squashfs data failed validation"
		puts "squashfs is #{fw.squashfs.size} bytes"
	else
		puts "No squashfs in this firmware file"
	end

	puts "Product ID: #{fw.product_id}"
	puts "Custom ID: #{fw.custom_id}"
	puts "Model ID: #{fw.model_id}"
	puts "Compat ID: #{fw.compat_id}"

	if optargs[:kernel]
		fw.kernel.write(optargs[:kernel])
		puts "Kernel data written to #{optargs[:kernel]}"
	end

	if optargs[:initrd]
		fw.initrd.write(optargs[:initrd])
		puts "initrd data written to #{optargs[:initrd]}"
	end

	if optargs[:defaults]
		if fw.defaults
			fw.defaults.write(optargs[:defaults])
			puts "defaults.tar.gz written to #{optargs[:defaults]}"
		else
			$stderr.puts "This firmware file does not have a defaults.tar.gz; not written"
		end
	end

	if optargs[:squashfs]
		if fw.squashfs
			fw.squashfs.write(optargs[:squashfs])
			puts "squashfs filesystem image written to #{optargs[:squashfs]}"
		else
			$stderr.puts "This firmware file does not have a squashfs filesystem image; not written"
		end
	end
end

if $0 == __FILE__
	case File.basename($0)
		when 'mkdns323fw'   then mkdns323fw(ARGV)
		when 'splitdns323fw' then splitdns323fw(ARGV)
		else
			$stderr.puts "Please call me as either 'mkdns323fw' or 'splitdns323fw'; symlinks are good"
			exit 1
	end
end