#!/usr/bin/env python
"""
#   Copyright (C) 2010  Michael Buesch <m@bues.ch>
#
#   This program is free software; you can redistribute it and/or modify
#   it under the terms of the GNU General Public License version 2
#   as published by the Free Software Foundation.
#
#   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.
"""

import sys
import getopt

def indexToName(index):
	D11UCODE_NAMETAG_START		= 0
	D11LCN0BSINITVALS24		= 1
	D11LCN0INITVALS24		= 2
	D11LCN1BSINITVALS24		= 3
	D11LCN1INITVALS24		= 4
	D11LCN2BSINITVALS24		= 5
	D11LCN2INITVALS24		= 6
	D11N0ABSINITVALS16		= 7
	D11N0BSINITVALS16		= 8
	D11N0INITVALS16			= 9
	D11UCODE_OVERSIGHT16_MIMO	= 10
	D11UCODE_OVERSIGHT16_MIMOSZ	= 11
	D11UCODE_OVERSIGHT24_LCN	= 12
	D11UCODE_OVERSIGHT24_LCNSZ	= 13
	D11UCODE_OVERSIGHT_BOMMAJOR	= 14
	D11UCODE_OVERSIGHT_BOMMINOR	= 15

	namemap = {
		D11UCODE_NAMETAG_START		: "start",
		D11LCN0BSINITVALS24		: "LCN0 bs initvals 24",
		D11LCN0INITVALS24		: "LCN0 initvals 24",
		D11LCN1BSINITVALS24		: "LCN1 bs initvals 24",
		D11LCN1INITVALS24		: "LCN1 initvals 24",
		D11LCN2BSINITVALS24		: "LCN2 bs initvals 24",
		D11LCN2INITVALS24		: "LCN2 initvals 24",
		D11N0ABSINITVALS16		: "N0A bs initvals 16",
		D11N0BSINITVALS16		: "N0 bs initvals 16",
		D11N0INITVALS16			: "N0 initvals 16",
		D11UCODE_OVERSIGHT16_MIMO	: "microcode 16 MIMO",
		D11UCODE_OVERSIGHT16_MIMOSZ	: "microcode 16 MIMO size",
		D11UCODE_OVERSIGHT24_LCN	: "microcode 24 LCN",
		D11UCODE_OVERSIGHT24_LCNSZ	: "microcode 24 LCN size",
		D11UCODE_OVERSIGHT_BOMMAJOR	: "bom major",
		D11UCODE_OVERSIGHT_BOMMINOR	: "bom minor",
	}
	try:
		return namemap[index]
	except KeyError:
		return "Unknown"

def parseHeader(hdr_data, sortByOffset):
	sections = []
	for i in range(0, len(hdr_data), 3 * 4):
		offset = ord(hdr_data[i + 0]) | (ord(hdr_data[i + 1]) << 8) |\
			(ord(hdr_data[i + 2]) << 16) | (ord(hdr_data[i + 3]) << 24)
		length = ord(hdr_data[i + 4]) | (ord(hdr_data[i + 5]) << 8) |\
			(ord(hdr_data[i + 6]) << 16) | (ord(hdr_data[i + 7]) << 24)
		index = ord(hdr_data[i + 8]) | (ord(hdr_data[i + 9]) << 8) |\
			(ord(hdr_data[i + 10]) << 16) | (ord(hdr_data[i + 11]) << 24)

		sections.append( (offset, length, index) )
	if sortByOffset:
		sections.sort(key = lambda x: x[0]) # Sort by offset
	else:
		sections.sort(key = lambda x: x[2]) # Sort by index
	return sections

def generateHeaderData(sections):
	data = []
	for section in sections:
		(offset, length, index) = section
		data.append(chr(offset & 0xFF))
		data.append(chr((offset >> 8) & 0xFF))
		data.append(chr((offset >> 16) & 0xFF))
		data.append(chr((offset >> 24) & 0xFF))
		data.append(chr(length & 0xFF))
		data.append(chr((length >> 8) & 0xFF))
		data.append(chr((length >> 16) & 0xFF))
		data.append(chr((length >> 24) & 0xFF))
		data.append(chr(index & 0xFF))
		data.append(chr((index >> 8) & 0xFF))
		data.append(chr((index >> 16) & 0xFF))
		data.append(chr((index >> 24) & 0xFF))
	return "".join(data)

def getSectionByIndex(sections, searchIndex):
	for section in sections:
		(offset, length, index) = section
		if searchIndex == index:
			return section
	return None

def parseHeaderFile(hdr_filepath, sortByOffset=False):
	try:
		hdr_data = file(hdr_filepath, "rb").read()
	except (IOError), e:
		print "Failed to read header file: %s" % e.strerror
		return None
	if len(hdr_data) % (3 * 4) != 0:
		print "Invalid header file format"
		return None
	return parseHeader(hdr_data, sortByOffset)

def dumpInfo(hdr_filepath):
	sections = parseHeaderFile(hdr_filepath)
	if not sections:
		return 1
	for section in sections:
		(offset, length, index) = section
		print "Index %2d   %24s    ==>  offset:0x%08X  length:0x%08X" %\
			(index, indexToName(index), offset, length)
	return 0

def extractSection(hdr_filepath, bin_filepath, extractIndex, outfilePath):
	sections = parseHeaderFile(hdr_filepath)
	if not sections:
		return 1
	section = getSectionByIndex(sections, extractIndex)
	if not section:
		print "Did not find a section with index %d" % extractIndex
		return 1
	(offset, length, index) = section
	try:
		bin_data = file(bin_filepath, "rb").read()
	except (IOError), e:
		print "Failed to read bin file: %s" % e.strerror
		return 1
	try:
		outfile = file(outfilePath, "wb")
		outfile.write(bin_data[offset : offset + length])
	except IndexError:
		print "Binfile index error."
		return 1
	except (IOError), e:
		print "Failed to write output file: %s" % e.strerror
		return 1
	return 0

def mergeSection(hdr_filepath, bin_filepath, mergeIndex, mergefilePath):
	sections = parseHeaderFile(hdr_filepath, sortByOffset=True)
	if not sections:
		return 1
	try:
		bin_data = file(bin_filepath, "rb").read()
	except (IOError), e:
		print "Failed to read bin file: %s" % e.strerror
		return 1
	try:
		merge_data = file(mergefilePath, "rb").read()
	except (IOError), e:
		print "Failed to open merge output file: %s" % e.strerror
		return 1
	newBin = []
	newSections = []
	newOffset = 0
	foundIt = False
	for section in sections:
		(offset, length, index) = section
		if index == mergeIndex:
			if foundIt:
				print "Confused. Multiple sections with index %d?" % index
				return 1
			foundIt = True
			# We overwrite this section
			newBin.append(merge_data)
			newSections.append( (newOffset, len(merge_data), index) )
			newOffset += len(merge_data)
		else:
			try:
				newBin.append(bin_data[offset : offset + length])
			except IndexError:
				print "Failed to read input data"
				return 1
			newSections.append( (newOffset, length, index) )
			newOffset += length
	if not foundIt:
		print "Did not find section with index %d" % mergeIndex
		return 1
	newBin = "".join(newBin)
	newHdr = generateHeaderData(newSections)
	try:
		file(bin_filepath, "wb").write(newBin)
		file(hdr_filepath, "wb").write(newHdr)
	except (IOError), e:
		print "Failed to write bin or header file: %s" % e.strerror
		return 1
	return 0

def usage():
	print "BRCM80211 firmware converter tool"
	print ""
	print "Usage: %s [OPTIONS]" % sys.argv[0]
	print ""
	print "  -H|--header FILE         Use FILE as header file"
	print "  -B|--bin FILE            Use FILE as bin file"
	print ""
	print "Actions:"
	print "  -d|--dump                Dump general information"
	print "  -x|--extract INDEX:FILE  Extract the section with index INDEX to FILE"
	print "  -m|--merge INDEX:FILE    Merges FILE into the bin file stream at INDEX"
	print "                           A Merge modifies the files specified in -H and -B"
	print ""
	print "  -h|--help                Print this help text"

def main():
	opt_header = None
	opt_bin = None
	opt_action = None
	opt_index = None
	opt_outfile = None
	opt_mergefile = None

	try:
		(opts, args) = getopt.getopt(sys.argv[1:],
			"hH:B:dx:m:",
			[ "help", "header=", "bin=", "dump", "extract=", "merge=", ])
	except getopt.GetoptError:
		usage()
		return 1
	for (o, v) in opts:
		if o in ("-h", "--help"):
			usage()
			return 0;
		if o in ("-H", "--header"):
			opt_header = v
		if o in ("-B", "--bin"):
			opt_bin = v
		if o in ("-d", "--dump"):
			opt_action = "dump"
		if o in ("-x", "--extract"):
			opt_action = "extract"
			try:
				v = v.split(':')
				opt_index = int(v[0])
				opt_outfile = v[1]
			except IndexError, ValueError:
				print "Invalid -x|--extract index number\n"
				usage()
				return 1
		if o in ("-m", "--merge"):
			opt_action = "merge"
			try:
				v = v.split(':')
				opt_index = int(v[0])
				opt_mergefile = v[1]
			except IndexError, ValueError:
				print "Invalid -m|--merge index and/or merge file name\n"
				usage()
				return 1
	if not opt_action:
		print "No action specified\n"
		usage()
		return 1
	if opt_action == "dump":
		if not opt_header:
			print "No header file specified\n"
			usage()
			return 1
		return dumpInfo(opt_header)
	elif opt_action == "extract":
		if not opt_header or not opt_bin:
			print "No header or bin file specified\n"
			usage()
			return 1
		return extractSection(opt_header, opt_bin, opt_index, opt_outfile)
	elif opt_action == "merge":
		if not opt_header or not opt_bin:
			print "No header or bin file specified\n"
			usage()
			return 1
		return mergeSection(opt_header, opt_bin, opt_index, opt_mergefile)
	return 1

if __name__ == "__main__":
	sys.exit(main())