Do memory diffing on specific block sizes

memdiff.py

import struct
import sys

def diffing(oldbuff, newbuff, op):
    offsets = set()

    for i in range(0, min(len(oldbuff), len(newbuff))):
        if op(oldbuff[i], newbuff[i]):
            offsets.add(i)

    return offsets

def alignbuff(buff, blksize):
    """ Aligns a buffer on a blocksize and returns a list of values """
    pattern = {
        1: 'B',
        2: 'H',
        4: 'I'
    }
    
    if len(buff) % blksize != 0:
        buff =  buff[:-(len(buff) % blksize)]

    count = len(buff) // blksize

    fmt = '<' + str(count) + pattern[blksize]

    return struct.unpack(fmt, buff)

def main():
    if len(sys.argv) < 5:
        print("usage: memdiff (8/16/32) 'pattern' file1 file2 ...")
        return
    
    size = sys.argv[1]
    pattern = sys.argv[2]
    files = sys.argv[3:]

    if size not in ["8", "16", "32"]:
        print("Invalid size: {}".format(size))
        return

    bytesize = int(size) // 8

    if len(pattern) != len(files)-1:
        print("Pattern has a wrong size of {}. Should be {}".format(len(pattern), len(files)-1))
        return

    cmpmap = {
            '+': lambda o, n: o < n,
            '-': lambda o, n: o > n,
            '=': lambda o, n: o == n
    }

    cur_offsets = None # Contains the currently active values after the search
    offsets_lst = {} # Contains the offset, and the chain of changed values
    
    for i, p in enumerate(pattern):
        if p not in cmpmap:
            print("invalid pattern char '{}'. Should be either '+' or '-'".format(p))
            return
        
        oldfile = files[i]
        newfile = files[i+1]

        oldbytes = open(oldfile, "rb").read()
        oldbytes = alignbuff(oldbytes, bytesize)

        newbytes = open(newfile, "rb").read()
        newbytes = alignbuff(newbytes, bytesize)
        
        offsets = diffing(oldbytes, newbytes, cmpmap[p])
        
        if cur_offsets is None:
            # When we have the first run we simply copy offsets
            cur_offsets = set()

            for o in offsets:
                cur_offsets.add(o)
                offsets_lst[o] = [oldbytes[o]]
        else:
            # If we already have existing offsets we do the union
            cur_offsets = cur_offsets.intersection(offsets)

        # Now we fill our offset_lst with the new entries
        for o in cur_offsets:
            newe = newbytes[o]
            offsets_lst[o].append(newe)

    for o in cur_offsets:
        values = " -> ".join(map(lambda a: "0x{:08x}".format(a), offsets_lst[o]))
        print("0x{:08x} : {}".format(o*bytesize, values))


if __name__ == '__main__': 
    main()