import sys, os, struct

def main():
    if len(sys.argv) != 2:
        print "USAGE: %s <TRUSTLET_DIRECTORY>" % sys.argv[0]
        return
    trustlet_dir = sys.argv[1]

    moduli = set()
    moduli_map = {}
    for root, dirs, files in os.walk(trustlet_dir):
        for name in files:
            if not (name.endswith(".tlbin") or name.endswith(".drbin") or name.endswith(".tabin")):
                continue
            
            print "----------"
            print "trustlet %s" % name
            print "dir %s" % root
            print 
            
            #Reading the length fields from the header
            trustlet = open(os.path.join(root, name), "rb").read()
            text_length = struct.unpack("<I", trustlet[0x34:0x38])[0]
            data_length = struct.unpack("<I", trustlet[0x3C:0x40])[0]
            off = text_length + data_length

            if off + 4 >= len(trustlet):
                print "Invalid offset, skipping."
                continue

            modulus_length = struct.unpack("<I", trustlet[off:off+4])[0]
            modulus = trustlet[off+4:off+4+modulus_length]
            moduli.add(modulus)
            print int(modulus.encode("hex"), 16)

            if modulus not in moduli_map:
                moduli_map[modulus] = set()
            moduli_map[modulus].add(root)
            
    print "------------------"
    print "moduli:"
    for N in moduli:
        print N.encode("hex")

    print "------------------"
    for N in moduli_map:
        print "modulus: %s -> %s" % (N.encode("hex"), moduli_map[N])

if __name__ == "__main__":
    main()