#!/usr/bin/python

import tarfile, gzip, array, hashlib, sys, struct
from binascii import hexlify
import os, subprocess
import random, tempfile
import getopt

## For Signed Images Only

BASE_CMD = "openssl dgst -sha256 "
CMD_SIGN = BASE_CMD + "-sign %(privkey)s -out %(outfile)s %(infile)s"
KINDLE_HACK_DIR = "/etc/uks"
KINDLE_HACK_KEYNAME = "pubprodkey01.pem"
SIGN_KEY = """
-----BEGIN RSA PRIVATE KEY-----
MIICXgIBAAKBgQDJn1jWU+xxVv/eRKfCPR9e47lPWN2rH33z9QbfnqmCxBRLP6mM
jGy6APyycQXg3nPi5fcb75alZo+Oh012HpMe9LnpeEgloIdm1E4LOsyrz4kttQtG
RlzCErmBGt6+cAVEV86y2phOJ3mLk0Ek9UQXbIUfrvyJnS2MKLG2cczjlQIDAQAB
AoGASLym1POD2kOznSERkF5yoc3vvXNmzORYkRk1eJkJuDY6yAbYiO7kDppqj4l8
wGogTpv98OMXauY8JgQj6tgO5LkY2upttukDr8uhE2z9Dh7HMZV/rDYa+9rybJus
RiAQDmF+VCzY2HirjpsSzgRu0r82NC8znNm2eGORys9BvmECQQDoIokOr0fYz3UT
SbHfD3engXFPZ+JaJqU8xayR7C+Gp5I0CgSnCDTQVgdkVGbPuLVYiWDIcEaxjvVr
hXYt2Ac9AkEA3lnERgg0RmWBC3K8toCyfDvr8eXao+xgUJ3lNWbqS0HtwxczwnIE
H49IIDojbTnLUr3OitFMZuaJuT2MtWzTOQJBAK6GCHU54tJmZqbxqQEDJ/qPnxkM
CWmt1F00YOH0qGacZZcqUQUjblGT3EraCdHyFKVT46fOgdfMm0cTOB6PZCECQQDI
s5Zq8HTfJjg5MTQOOFTjtuLe0m9sj6zQl/WRInhRvgzzkDn0Rh5armaYUGIx8X0K
DrIks4+XQnkGb/xWtwhhAkEA3FdnrsFiCNNJhvit2aTmtLzXxU46K+sV6NIY1tEJ
G+RFzLRwO4IFDY4a/dooh1Yh1iFFGjcmpqza6tRutaw8zA==
-----END RSA PRIVATE KEY-----
"""

NEW_KEY = """
-----BEGIN PUBLIC KEY-----
MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQDJn1jWU+xxVv/eRKfCPR9e47lP
WN2rH33z9QbfnqmCxBRLP6mMjGy6APyycQXg3nPi5fcb75alZo+Oh012HpMe9Lnp
eEgloIdm1E4LOsyrz4kttQtGRlzCErmBGt6+cAVEV86y2phOJ3mLk0Ek9UQXbIUf
rvyJnS2MKLG2cczjlQIDAQAB
-----END PUBLIC KEY-----
"""

INSTALL_SCRIPT = """
#!/bin/sh

_FUNCTIONS=/etc/rc.d/functions
[ -f ${_FUNCTIONS} ] && . ${_FUNCTIONS}

. /etc/sysconfig/mntus

#Restore original Kindle signing key
cat > /etc/uks/pubprodkey01.pem <<EOF
-----BEGIN PUBLIC KEY-----
MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCxfpiZ1dbdSOgrikqXD6lESUrD
5l52nN50iMh2vDcmW/FzkPDv0eRf1ci6w3ifhmHwqDK9OYNnowPapzUHAiHukXjW
rOC3fZYzgAxzIPN4NyUw369zFK2AALZnXptc68D/xxtZ94porf+kLtw/4vF2NhHs
XtchrpvID+Jhkor8MQIDAQAB
-----END PUBLIC KEY-----
EOF

update_progressbar 100

return 0
"""

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 = hashlib.md5()
  m.update(s)
  return hexlify(m.digest())

def runCommand(cmd):
    """wrapper to simplify the execution of external programs.
    @param cmd: Command line to be executed
    @type cmd: string

    @return: tuple (exit code, stdout, stderr).
    @rtype: tuple
    """
    ssl = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    status = ssl.wait()
    out, err = ssl.communicate()
    return status, out, err

def convert_bin(fromVer, toVer, kid, type, input, output):
    print "making bin file"
    BLOCK_SIZE=64
    if type=="FB02" or type=="FB01":
      BLOCK_SIZE=131072

    f = open(input, "rb").read()
    of = open(output,"wb")
    #C4 1D 3C 07 C2 B5 A0 08
    header = struct.pack("<4sIIHBB", type, int(fromVer), int(toVer), int(kid), 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))
    of.close()
    print "output written to "+output

def usage():
  print """
    Kindle Firmware Update tool Lite v0.11
    Copyright (c) 2009-2010 Igor Skochinsky & Jean-Yves Avenard
    Usage:

    kindle_update_tool.py c fromVersion toVersion deviceCode updateType inputTgz outputBin
      Convert a GZIPPED TAR file into a firmware update file.
"""

if len(sys.argv)<5:
    usage()
elif sys.argv[1]=="c":
    fromVer = sys.argv[2]
    toVer = sys.argv[3]
    kid = sys.argv[4]
    type = sys.argv[5]
    input = sys.argv[6]
    output = sys.argv[7]
    convert_bin(fromVer, toVer, kid, type, input, output)
else:
    usage()
    sys.exit(2)
