MC Anvil region unpacking/packing script

region.py

import sys, os, zlib, struct

def decompress_chunk(chunk):
    l, = struct.unpack_from(">I", chunk, 0)
    assert l <= len(chunk)-4, "Invalid length"
    assert chunk[4:5] == b"\x02", "Only zlib-compressed chunks are supported"
    return zlib.decompress(chunk[5:4+l])

def compress_chunk(nbt):
    chunk = b"\x02" + zlib.compress(nbt)
    return struct.pack(">I", len(chunk)) + chunk

def deserialize_region(data):
    if len(data) == 0: return # empty region, it happens
    locs = struct.unpack_from(">"+"I"*(32*32), data, 0)
    tss = struct.unpack_from(">"+"I"*(32*32), data, 4096)
    for i, (loc, ts) in enumerate(zip(locs, tss)):
        z, x = divmod(i, 32)
        if loc:
            offs = loc>>8<<12
            clen = (loc&0xff)<<12
            yield x, z, ts, decompress_chunk(data[offs:offs+clen])

def serialize_region(entries):
    assert len(entries) == 1024
    region = bytearray()
    region += b"\x00" * 8192
    for i, e in enumerate(entries):
        if e is None: continue
        nbt, ts = e
        chunk = compress_chunk(nbt)
        clen = (len(chunk) + (1<<12) - 1) >> 12
        assert clen < 256, "Compressed chunk too long for Anvil format"
        chunk += b"\x00" * ((clen<<12) - len(chunk))
        offs = len(region)>>12
        region += chunk
        struct.pack_into(">I", region, 4*i, (offs<<8) | clen)
        struct.pack_into(">I", region, 4096 + 4*i, int(ts))
    return bytes(region)

def usage():
    print("region.py unpack <region x> <region z>")
    print("region.py pack <region x> <region z>")
    print("This script should be ran within a world folder")
    exit(2)

if __name__ == "__main__":
    args = sys.argv[1:]
    if not args: usage()
    if args[0] == "unpack" and len(args) == 3:
        rx, rz = int(args[1]), int(args[2])
        ox, oz = rx*32, rz*32
        base = "region/r.%d.%d" % (rx, rz)
        with open(base+".mca","rb") as regf:
            reg = regf.read()
        if not os.path.isdir(base):
            os.mkdir(base)
        for x, z, ts, nbt in deserialize_region(reg):
            cpath = base+"/chunk.%d.%d.nbt" % (ox+x, oz+z)
            with open(cpath, "wb") as chf:
                chf.write(nbt)
            os.utime(cpath, (ts, ts))
            print("Unpacked %s, %d bytes uncompressed" % (cpath, len(nbt)))
    elif args[0] == "pack" and len(args) == 3:
        rx, rz = int(args[1]), int(args[2])
        base = "region/r.%d.%d" % (rx, rz)
        if not os.path.isdir(base):
            print("Could not find chunk directory %s/" % base)
            exit(1)
        ox, oz = rx*32, rz*32
        entries = [None] * 1024
        for dz in range(32):
            for dx in range(32):
                cpath = base + "/chunk.%d.%d.nbt" % (ox+dx, oz+dz)
                try:
                    st = os.stat(cpath)
                except FileNotFoundError:
                    continue
                with open(cpath, "rb") as f:
                    nbt = f.read()
                entries[dz*32+dx] = (nbt, st.st_mtime)
        mca = serialize_region(entries)
        print("New region size: %d bytes" % len(mca))
        os.rename(base+".mca", base+".mca.old")
        with open(base+".mca", "wb") as regf:
            regf.write(mca)
        print("Comparing to old region data")
        with open(base+".mca.old","rb") as regf:
            oreg = regf.read()
        newd = {}
        oldd = {}
        for x, z, ts, nbt in deserialize_region(mca):
            newd[(ox+x, oz+z)] = (ts, nbt)
        for x, z, ts, nbt in deserialize_region(oreg):
            oldd[(ox+x, oz+z)] = (ts, nbt)
        if newd == oldd: print("Regions contain identical data")
        else:
            for k in set(newd.keys()) | set(oldd.keys()):
                if k in newd and k in oldd:
                    if newd[k] != oldd[k]: print("Chunk %r modified, (%d bytes, timestamp %d) -> (%d bytes, timestamp %d)" % (k, len(oldd[k][1]), oldd[k][0], len(newd[k][1]), newd[k][0]))
                else:
                    assert k in newd or k in oldd
                    if k in newd:
                        print("Added new chunk %r" % (k,))
                    if k in oldd:
                        print("Removed chunk %r" % (k,))
    else:
        usage()