#!/usr/bin/python
# Kindle Firmware Update tool v0.1 Copyright (c) 2009 Igor Skochinsky
# History:
#  2009-03-10 Initial release

import tarfile, gzip, array, md5, sys, struct
from binascii import hexlify

if sys.hexversion >= 0x3000000:
  print "This script is incompatible with Python 3.x. Please install Python 2.6.x from python.org"
  sys.exit(2)

def dm(s):
  arr = array.array('B',s)
  for i in xrange(len(arr)):
    b = arr[i]^0x7A
    arr[i] = (b>>4 | b<<4)&0xFF
  return arr.tostring()

def md(s): #opposite of dm
  arr = array.array('B',s)
  for i in xrange(len(arr)):
    b = arr[i]
    b = (b>>4 | b<<4)&0xFF
    arr[i] = b^0x7A
  return arr.tostring()

def s_md5(s): #return md5 in string format
  m = md5.new()
  m.update(s)
  return hexlify(m.digest())

def extract_bin(binname):
    f = file(binname, "rb")
    sig, fromVer, toVer, devCode, optional = struct.unpack("<4sIIHBx", f.read(16))
    
    if sig=="FC02":
      typ="OTA update"
    elif sig=="FB01":
      typ="Manual update"
    else:
      print "Not a Kindle update file!"
      return

    print "Signature: %s (%s)"%(sig,typ)
    if sig=="FC02":
      print "min version: %d"%(fromVer)
      print "max version: %d"%(toVer)
      print "device code: %d%d"%divmod(devCode,256)
      print "optional: "+("yes" if optional else "no")
    print "md5 of tgz: %s"%dm(f.read(32))
    if sig=="FC02":
    	f.seek(64)
    elif sig=="FB01":
        f.seek(131072)
    file(binname+".tgz","wb").write(dm(f.read()))

def make_bin(basename, filelist, type, kver):
    tgz_fname = "temp.tar.gz"
    tar = tarfile.open(tgz_fname,"w:gz")

    dat_list = ""
    for name in filelist:
      print "adding %s"%name
      tarinfo = tar.gettarinfo(name)
      tarinfo.uid = tarinfo.gid = 0
      tarinfo.uname = tarinfo.gname = "root"
      if name.endswith(".sh"):
        tarinfo.mode = 0100755 #rwxr-xr-x
        fid = 129
      else:
        tarinfo.mode = 0100644 #rw-r--r--
        fid = 128
      tarinfo.type = tarfile.REGTYPE
      inf = file(name,"rb")
      tar.addfile(tarinfo, inf)
      inf.seek(0)
      #129 a2592fe6898d468fd64f00c5b4b04ad7 test.sh 0 Test_script
      dat_list+="%d %s %s 0 %s\n"%(fid, s_md5(inf.read()), name, name+"_file")

    file(basename+".dat","wb").write(dat_list)
    tar.add(basename+".dat")
    tar.close()

    print "making bin file"
    if type==2:
      BLOCK_SIZE=64
      sig = "FC02"
    else:
      BLOCK_SIZE=131072
      sig = "FB01"

    f = file(tgz_fname, "rb").read()
    of = file(basename+".bin","wb")
    #C4 1D 3C 07 C2 B5 A0 08
    header = struct.pack("<4sIIHBB", sig, 0x0, 0x7fffffff, kver, 0, 0x13) #signature, fromVer, toVer, devCode, optional
    of.write(header)
    of.write(md(s_md5(f)))
    of.write("\0"*(BLOCK_SIZE - of.tell()))
    of.write(md(f))
    print "output written to "+basename+".bin"

def usage():
  print """Usage:
    kindle_update_tool.py e update_mmm.bin
      Extract a Kindle or Kindle 2 firmware update file. Outputs a .tgz file with decrypted content.
    kindle_update_tool.py m [-k2] name file1 [file2 ...]
      Makes a Kindle or Kindle 2 (if -k2 specified) OTA firmware update file from the list of files.
      "name" is the update file suffix (final file will be called update_name.bin).
      Any file with .sh extension will be marked as a shell script to be executed."""
  
print "Kindle Firmware Update tool v0.1 Copyright (c) 2009 Igor Skochinsky"

if len(sys.argv)<3:
  usage()
elif sys.argv[1]=="e":
  extract_bin(sys.argv[2])
elif sys.argv[1]=="m":
  kver = 1
  if sys.argv[2]=="-k2": 
    kver = 2
  if sys.argv[2]=="-k3": 
    kver = 2
  if sys.argv[2]=="-k4": 
    kver = 2
  name = sys.argv[kver+1]
  filelist = sys.argv[kver+2:]
  if sys.argv[2]=="-k3": 
    kver = 4
  if sys.argv[2]=="-k4": 
    kver = 3
  make_bin("update_"+name, filelist, 2, kver)
else:
  usage()