#!/usr/local/bin/python
###############################################################################
###
### (c) 2012 by Ollopa / www.isysop.com
###
### Python script to apply a copy-patch to a file.
###
###
import os, sys, hashlib, binascii

global origMd5, finalMd5, rdh

### State types
class ST:
    HEADER = 0
    CPP = 1
    BINP = 2

if len(sys.argv) < 3:
    print "Error:  Too few arguments!"
    print "Usage:", sys.argv[0], "<binary file to patch> <source file(s)> ... < patchFile"
    print "For example, to apply patch 'update' to file 'image' using files" 
    print "'a' and 'b' as source files:"
    print sys.argv[0], "image a b < update"
    exit(1)

fileList = dict()

### Calculate unmodified ramdisk md5 sum
rdh = open(sys.argv[1], 'r+b')
origMd5 = hashlib.md5(rdh.read()).hexdigest().lower()

finalMd5 = ""

### Build a dictionary of md5 sums <-> filenames
for filename in sys.argv[2:]:
    fh = open(filename, 'rb')
    md5 = hashlib.md5(fh.read()).hexdigest()
    fh.close()
    fileList[md5] = filename

state = ST.HEADER
expectedseq = 0
### Read patch line-by-line from stdin
for line in sys.stdin.readlines():
    line = line.strip()

    ### Ignore blank lines and comments
    if line == "" or line[0] == "#":
        continue

    words = line.split()
    if state == ST.HEADER:
        if words[0] == "@-":
            ### Compare unmodified ramdisk md5 with md5 in patch file
            if origMd5 != words[1].lower():
                print "Error: Starting ramdisk md5 does not match!"
                exit(1)
        elif words[0] == "@+":
            ### Store modified ramdisk md5 for comparison after patching
            finalMd5 = words[1].lower()
        elif words[0].upper() == "CPP:":
            ### Attempt to open file for copy-patching
            if len(words) != 3:
                print "Error: Invalid CPP line format:", line
                exit(1)
            try:
                filename = fileList[words[2]]
            except:
                print "Error: Could not find file for", words[1]
                exit(1)
            ### print "Using", filename, "for", words[1]
            sys.stdout.write("Patching-in \"" + filename+ "\" ")
            fh = open(filename, 'rb')
            expectedseq = 1
            state = ST.CPP
        elif words[0].upper() == "BINP:":
            ### Invalid
            print "Parse Error: BINP block not allowed before CPP block!"
            exit(1)
        else:
            print "Parse Error: Unexpected input!"
            print "Offending line:", line
            exit(1)
    elif state == ST.CPP:
        if words[0] == "@-" or words[0] == "@+":
            ### Invalid state
            print "Parse Error: @- or @+ in CPP block!"
            exit(1)
        elif words[0].upper() == "CPP:":
            print "Done."
            ### Close old file and attempt to open new file
            if len(words) != 3:
                print "Error: Invalid CPP line format:", line
                exit(1)
            try:
                filename = fileList[words[2]]
            except:
                print "Error: Could not find file for", words[1]
                exit(1)
            ### print "Using", filename, "for", words[1]
            sys.stdout.write("Patching-in \"" + filename+ "\" ") 
            fh.close()
            fh = open(filename, 'rb')
            expectedseq = 1
        elif words[0].upper() == "BINP:":
            ### close file and move to binary patch statea
            fh.close()
            print "Done."
            print
            print "Performing binary patches to file system structure ",
            state = ST.BINP
        else:
            ### Perform copy-patch
            ### Line formate is:
            ### seq: offset ... len
            ### Seq is the read sequence in the source file
            ### offset ... is a list of locations in the ramdisk file to patch
            ### len is the length of data to read from source (block size)
            if words[0][-1] == ":":
                seq = int(words[0][:-1])
                length = int(words[-1])
                block = fh.read(length)
                if seq != expectedseq:
                    print "Error: Unexpected sequence value."
                    print "Expected", expectedseq, "but got", seq
                    print "Offending line:", line
                    exit(1)
                for location in words[1:-1]:
                    if location.upper() == "SKIP":
                        continue
                    rdh.seek(int(location))
                    rdh.write(block)
                    sys.stdout.write(".")
                expectedseq += 1
            else:
                print "Parse Error: malformed copy-patch line!"
                print "Offending line:", line
                exit(1)
    elif state == ST.BINP:
        if words[0] == "@-" or words[0] == "@+":
            print "Parse Error: @- or @+ in BINP block!"
            exit(1)
        elif words[0].upper() == "CPP:":
            print "Parse Error: CPP block not allowed after BINP block!"
            exit(1)
        else:
            if words[0][-1] == ":" and len(words) == 2:
                location = int(words[0][:-1])
                hexdata = words[1]
                ### Even numbers only
                if len(hexdata) & 1:
                    print "Parse Error: malformed binary-patch line!"
                    print "Hex data length is not even"
                    print "Offending line:", line
                    exit(1)
                rdh.seek(location)
                rdh.write(binascii.unhexlify(hexdata))
                sys.stdout.write(".")
            else:
                print "Parse Error: malformed binary-patch line!"
                print "Offending line:", line
                exit(1)
print
### Compare final md5 sum with value in patch file
rdh.seek(0)
if finalMd5 != hashlib.md5(rdh.read()).hexdigest().lower():
    print "Error: Ending ramdisk md5 does not match!"
    exit(1)
else:
    print "Patch successful!  md5:", finalMd5