#!/usr/bin/lua
-- Customize an app – step 2: modify ihex

assert(tonumber('AABB', 16) == 43707, "Endianness problem!")

local bit = require("bit")

SEPARATOR = ',' -- Seperator used in modifications file


-- Extend string library with trim(str) function
string.trim = function(s)
	return (s:gsub("^%s*(.-)%s*$", "%1"))
end

-- Extend string library with replace2(str, pos, rep) function where len(rep) = 2
string.replace2 = function(str, pos, rep)
	return string.sub(str, 1, pos - 1) .. rep .. string.sub(str, pos + 2)
end

-- Extend table library with dictsize(dict) function
table.dictsize = function(dict)
	local count = 0
	for _ in pairs(dict) do
		count = count + 1
	end
  return count
end


-- Representation of ihex entry
local function IHexEntry(slen, saddr, skind, sdata)
	local self = {
		slen = slen,
		saddr = saddr,
		skind = skind,
		sdata = sdata,
	}

	-- Is data entry?
	function self.isData()
		return self.skind == "00"
	end

	-- Is eof entry?
	function self.isEof()
		return self.skind == "01"
	end

	-- Is extended segment address entry?
	function self.isESA()
		return self.skind == "02"
	end

	-- Is extended linear address entry?
	function self.isELA()
		return self.skind == "04"
	end

	-- Address from the entry
	 function self.addr()
		return tonumber(self.saddr, 16)
	end

	-- Length from the entry
	function self.len()
		return tonumber(self.slen, 16)
	end

	-- Modify the entry (if it should be modified)
	function self.modify(base_addr, modifications)
		for i = 0, self.len() - 1 do
			local byte_addr = base_addr + self.addr() + i
			if (modifications[byte_addr]) then
				self.sdata = string.replace2(self.sdata, i*2 + 1, modifications[byte_addr])
				modifications[byte_addr] = nil -- The modification has been used
			end
		end 
	end

	-- Calculate checksum
	function self.checksum()
		local checksum = 0

		for i = 1, string.len(sdata), 2 do
			checksum = checksum + tonumber(string.sub(self.sdata, i, i+1), 16)
		end
		checksum = checksum + tonumber(self.slen, 16)
		checksum = checksum + tonumber(string.sub(self.saddr, 1, 2), 16)
		checksum = checksum + tonumber(string.sub(self.saddr, 3, 4), 16)
		checksum = checksum + tonumber(self.skind, 16)

		checksum = bit.band(checksum, 0xFF)
		checksum = bit.bnot(checksum) + 1
		checksum = bit.band(checksum, 0xFF)

		return checksum
	end

	-- Format the entry as ihex line
	function self.line()
		local checksum = self.checksum()

		return ":".. self.slen .. self.saddr .. self.skind .. self.sdata .. string.format("%02X", checksum)
	end

	-- Validate the entry
	function self.validate(checksum)
		if self.len() ~= string.len(sdata) / 2 then return false end
		if self.checksum() ~= checksum then return false end

		return true
	end

	return self
end


-- Parse line into ihex entry
local function parseIHex(line)
	local len, addr, kind, data, checksum
	local entry

	-- Parse line:
	len, addr, kind, data, checksum =
				string.match(string.trim(line), "^:(%x%x)(%x%x%x%x)(%x%x)(%x*)(%x%x)$")
	if len == nil then
		io.stderr:write(string.format("ERROR: Cannot parse ihex file line: %s\n", line))
		return nil
	end

	-- Validate entry:
	entry = IHexEntry(len, addr, kind, data)
	if not entry.validate(tonumber(checksum, 16)) then
		io.stderr:write(string.format("ERROR: Invalid ihex file line: %s\n", line))
		return nil
	end

	return entry
end


-- Process ihex applying modifications
local function applyModifications(in_file, out_file, modifications)
	local fin, fout
	local upper_addr = 0
	local eof = false
	local msg

	-- Open files:
	fin, msg = io.open(in_file, 'r')
	if fin == nil then
		io.stderr:write(string.format("ERROR: Cannot open ihex file: %s\n", msg))
		return nil
	end
	fout, msg = io.open(out_file, 'w')
	if fout == nil then
		io.stderr:write(string.format("ERROR: Cannot create new ihex file: %s\n", msg))
		fin:close()
		return nil
	end

	-- Process ihex file:
	while not eof do
		-- Read line:
		local line = fin:read("*l")
		if line == nil then
			io.stderr:write("ERROR: Unexpected end of the ihex file\n")
			fin:close()
			fout:close()
			return nil
		end

		-- Parse line:
		local entry = parseIHex(line)
		if entry == nil then
			fin:close()
			fout:close()
			return nil
		end

		-- Interpret entry:
		if entry.isData() then
			entry.modify(upper_addr, modifications)
		elseif entry.isEof() then
			eof = true
		elseif entry.isESA() then
			upper_addr = tonumber(entry.sdata, 16) * 16
		elseif entry.isELA() then 
			upper_addr = bit.lshift(tonumber(entry.sdata, 16), 16)
		else
			assert((entry.stype == "03") or (entry.stype == "05")) -- Do nothing
		end

		-- Write entry to modified ihex
		fout:write(entry.line(), '\n')
	end

	-- Close files:
	fin:close()
	fout:close()

	-- Were all modifications applied?
	if table.dictsize(modifications) ~= 0 then
		io.stderr:write("WARNING: Some modifications are not applied\n")
	end

	return true
end


-- Read modifications file
local function readModifications(modifications_file, me)
	local addresses = {}
	local modifications = {}
	local md5sum
	local was_me = false
	local line
	local fm
	local msg

	-- Open modifications file:
	fm, msg = io.open(modifications_file, 'r')
	if fm == nil then
		io.stderr:write(
			string.format("ERROR: Cannot open modifications file: %s\n", msg)
		)
		return nil, nil
	end

	-- md5sum:
	md5sum = fm:read("*l")

	-- Addresses:
	line = fm:read("*l")
	for str in string.gmatch(line, "([^"..SEPARATOR.."]+)") do
		table.insert(addresses, tonumber(str, 16))
	end

	-- Values:
	line = fm:read("*l")
	while (line ~= nil) and (not was_me) do
		local mote = string.match(line, "([^"..SEPARATOR.."]+)")
		if string.upper(mote) == string.upper(me) then
			local index = 1
			line = string.sub(line, string.len(mote) + 2)
			for str in string.gmatch(line, "([^"..SEPARATOR.."]+)") do
				assert(index <= #addresses)
				assert(string.len(str) == 2)
				modifications[addresses[index]] = str
				index = index + 1
			end
			assert(index == #addresses + 1)
			was_me = true
		end
		line = fm:read("*l")
	end

	-- Close modifications file
	fm:close()

	-- Are there modifications for me?
	if not was_me then
		io.stderr:write(
			string.format("WARNING: No modifications for me (%s)\n", me)
		)
	end

	return md5sum, modifications
end


-- Verify md5sum of ihex file
local function verifyIHex(ihex_file, md5sum)
	local output
	local new_md5sum
	local f

	f = io.popen("md5sum " .. ihex_file)
	output = f:read('*l')
	f:close()

	new_md5sum = string.match(output, "^(%w+)%s+[^%s]+$")
	if md5sum ~= new_md5sum then
		io.stderr:write("ERROR: The modifications file was prepared for a different ihex file\n")
		io.stderr:write(string.format(" Local md5sum: %s\n", new_md5sum))
		io.stderr:write(string.format(" Original md5sum: %s\n", md5sum))
		return nil
	end

	return true
end

-- Get this device name
local function getMyName()
	local name
	local mac
	local f
	local msg

	-- Open file:
	f, msg = io.open("/sys/class/net/eth0/address", 'r')
	if f == nil then
		io.stderr:write(string.format("ERROR: Cannot read my MAC address: %s\n", msg))
		return nil
	end

	-- Read MAC address:
	mac = f:read("*a")
	f:close()

	-- Get CherryMote name:
	name = string.sub(mac, 13, 14) .. string.sub(mac, 16, 16)

	return name
end


-- Get command line parameters
local function getArgs()
	if (#arg ~= 3) and (#arg ~= 4) then
		io.stderr:write(
			string.format("Usage: %s modifications_file ihex_file path_for_modified_ihex_file [mote_ID]\n", arg[0])
		)
		return nil, nil, nil, nil
	end

	-- mote ID: from command line parameter or this mote
	if (#arg == 4) then
		mote = arg[4]
	else
		mote = getMyName()
	end

	return arg[1], arg[2], arg[3], mote
end


-- Modify the ihex – main function
local function modify()
	-- Get command line parameters:
	local modifications_file, in_file, out_file, mote = getArgs()
	if modifications_file == nil then return 1 end

	-- Load modifications:
	local md5sum, modifications = readModifications(modifications_file, mote)
	if md5sum == nil then return 1 end

	-- Verify ihex file:
	ihex_ok = verifyIHex(in_file, md5sum)
	if ihex_ok == nil then return 1 end

	-- Process ihex applying modifications
	local status = applyModifications(in_file, out_file, modifications)
	if status == nil then return 1 end

	return 0
end


-- Run it!
os.exit(modify())
