mkkdz.py

#!/usr/bin/env python

"""
Copyright (C) 2016 Elliott Mitchell <ehem+android@m5p.com>

	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 3 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, see <http://www.gnu.org/licenses/>.
"""

from __future__ import print_function
import os
import sys
import argparse

# our tools are in "libexec"
sys.path.append(os.path.join(sys.path[0], "libexec"))

import kdz


class KDZFileTools(kdz.KDZFile):
	"""
	LGE KDZ File creation tool
	"""

	indir = "kdzextracted"

#kdz.KDZFile._dz_header


	def loadParams(self):
		"""
		Loads the .kdz.params file for output, saves the values
		"""

		params = dict()
		file = open(os.path.join(self.indir, ".kdz.params"), "rt")
		line = file.readline()
		while len(line) > 0:
			line.lstrip()
			line = line.partition("#")[0]
			if len(line) == 0:
				line = file.readline()
				continue
			parts = line.partition("=")
			if len(parts[1]) == 0:
				print("[!] Bad line in {:s}'s parameter file".format(self.indir), file=sys.stderr)
			var = parts[0].rstrip()
			try:
				val = int(parts[2].strip())
			except ValueError:
				val = parts[2].strip()
			params[var] = val
			line = file.readline()

		file.close()

		for k in 'version', 'dataStart', 'payload0', 'payload0head':
			if k not in params:
				print("Parameter value \"{:s}\" is missing, unable to continue".format(k))
				sys.exit(1)

		if params['version'] != 2:
			print("[!] File format version not 2, cannot continue", file=sys.stderr)
			sys.exit(1)

		self.dataStart = params['dataStart']

		i = 0
		self.headers = []
		headers = {}
		self.payload = []
		self.files = {}
		name = "payload"+str(i)
		while name in params:
			self.payload.append(params[name])
			headers[params[name+"head"]] = params[name]
			i += 1
			name = "payload"+str(i)
		for i in range(i):
			self.headers.append(headers[i])

	def cmdCreateFile(self):
		"""
		Create the specified KDZ file
		"""

		out = open(self.kdzfile, "wb")
		current = self.dataStart
		out.seek(current, os.SEEK_SET)

		for name in self.payload:
			print("[+] Writing {:s} to output file {:s}".format(name, self.kdzfile))
			inf = open(os.path.join(self.indir, name), "rb")
			buf = " "
			while len(buf) > 0:
				buf = inf.read(4096)
				out.write(buf)
			self.files[name] = [current, out.tell() - current]
			inf.close()
			current = out.tell()

		self.files[self.headers[-1]].append(0)

		print("\n[+] Writing headers to {:s}".format(self.kdzfile))

		out.seek(0, os.SEEK_SET)
		out.write(self._dz_header)
		for name in self.headers:
			# last record
			if len(self.files[name]) > 2:
				out.write(b'\x03')
			head = {
				'name':		name.encode("utf8"),
				'length':	self.files[name][1],
				'offset': 	self.files[name][0],
			}
			buf = self.packdict(head)
			out.write(buf)
		out.close()
		print("[+] Done!")

	def cmdList(self):
		pass


	def parseArgs(self):
		# Parse arguments
		parser = argparse.ArgumentParser(description='LG KDZ File creator by Elliott Mitchell')
		parser.add_argument('-f', '--file', help='KDZ File to read', action='store', required=True, dest='kdzfile')
		group = parser.add_mutually_exclusive_group(required=True)
		group.add_argument('-l', '--list', help='list partitions', action='store_true', dest='listOnly')
		group.add_argument('-m', '--make', help='extract all partitions', action='store_true', dest='createFile')
		parser.add_argument('-d', '--dir', help='input directory', action='store', dest='indir')

		return parser.parse_args()

	def main(self):
		args = self.parseArgs()
		self.kdzfile = args.kdzfile

		if args.indir:
			self.indir = args.indir

		self.loadParams()

		if args.listOnly:
			self.cmdList()

		elif args.createFile:
			self.cmdCreateFile()

if __name__ == "__main__":
	kdztools = KDZFileTools()
	kdztools.main()