converter for RamsesA's wow maps

convert.py

#!/usr/bin/python

import redstone as rs
import numpy
import multiprocessing
import sys
import os
import os.path

# amount of overlap between layers
OVERLAP = 64
# vanilla world height
HEIGHT = 256

def main():
	if len(sys.argv) != 3:
		print("usage: {} [base directory] [output]".format(sys.argv[0]), file=sys.stderr)
		return 1

	src = os.path.expanduser(sys.argv[1])
	dest = os.path.expanduser(sys.argv[2])
	name = os.path.split(src)[-1]
	base = os.path.join(*os.path.split(src)[:-1])

	# check name and base are good
	if not os.path.isdir(base) or name == '':
		print("invalid source: `{}'".format(src), file=sys.stderr)
		print("(expected worlds in `{}'".format(base), file=sys.stderr, end=" ")
		print("with names like {})".format(repr(name + "+001")), file=sys.stderr)
		return 1

	# make the destination directory
	destregions = os.path.join(dest, 'region')
	if not os.path.isdir(destregions):
		os.makedirs(destregions)

	print("converting {}".format(name))
	
	# scan the input directories for all possible levels
	levels = {}
	for dir in os.listdir(base):
		if not dir.startswith(name):
			continue
		if dir == name:
			levels[0] = os.path.join(base, dir)
		if '+' in dir:
			stem, num = dir.split('+', 1)
			if stem != name:
				continue
			try:
				num = int(num)
				levels[num] = os.path.join(base, dir)
			except ValueError:
				continue
		if '-' in dir:
			stem, num = dir.split('-', 1)
			if stem != name:
				continue
			try:
				num = int(num)
				levels[-num] = os.path.join(base, dir)
			except ValueError:
				continue

	# remap levels to be all greater than or equal to 0
	newlevels = {}
	offset = min(levels.keys())
	for i, path in levels.items():
		newlevels[i - offset] = path
	levels = newlevels

	sections_high = get_layer_offset(len(levels) - 1) + HEIGHT
	if sections_high % 16 != 0:
		print("calculated height is *not* a multiple of 16!", file=sys.stderr)
		return 1
	sections_high = sections_high // 16
	print("found {} distinct levels".format(len(levels)))
	print("resulting world will be {} chunk sections high".format(sections_high))

	# find out what mca files are relevant in each case
	regions = set()
	for path in levels.values():
		files = [e for e in os.listdir(os.path.join(path, 'region')) if e.endswith('.mca')]
		regions.update(files)

	print("found {} relevant regions".format(len(regions)))

	# convert them all
	pool = multiprocessing.Pool()
	map_args = [(i, len(regions), region, levels, dest) for (i, region) in enumerate(regions)]
	pool.imap_unordered(do_region_job, map_args)
	pool.close()
	pool.join()

	print("done")
	return 0

def do_region_job(args):
	i, total, region, levels, dest = args
	print("converting region {}/{}".format(i+1, total))
	convert_region(region, levels, dest)

def convert_region(name, srcmap, dest):
	# map integers to librs regions
	srcregions = {}
	for i, src in srcmap.items():
		fname = os.path.join(src, 'region', name)
		if os.path.exists(fname):
			srcregions[i] = rs.Region.open(fname)

	# open our destination region
	fname = os.path.join(dest, 'region', name)
	destregion = rs.Region.open(fname, True)

	# iterate through all available chunks
	for x in range(32):
		for z in range(32):
			srcchunks = {}
			for i, reg in srcregions.items():
				if reg.contains_chunk(x, z):
					srcchunks[i] = rs.NBT.parse_from_region(reg, x, z)

			# if we have data, convert it and write it out
			if len(srcchunks) == 0:
				continue
			destchunk = convert_chunk(srcchunks, len(srcmap))
			if destchunk:
				destchunk.write_to_region(destregion, x, z)

	# flush output
	destregion.flush()

# helper class to contain / render sections onto a huge array
class Section(object):
	def __init__(self, tag):
		self.y = tag['Y'].value
		blocks = tag['Blocks'].value
		blocklight = tag['BlockLight'].value
		skylight = tag['SkyLight'].value
		data = tag['Data'].value

		self.blocks = numpy.frombuffer(blocks, dtype=numpy.uint8).reshape((16, 16, 16))
		self.blocklight = self.parse_nibbles(blocklight)
		self.skylight = self.parse_nibbles(skylight)
		self.data = self.parse_nibbles(data)

	def parse_nibbles(self, s):
		dat = numpy.frombuffer(s, dtype=numpy.uint8).reshape((8, 16, 16))
		dat_expanded = numpy.empty((16, 16, 16), dtype=numpy.uint8)
		dat_expanded[::2,:,:] = dat & 0x0f
		dat_expanded[1::2,:,:] = (dat & 0xf0) >> 4
		return dat_expanded
		
	def render(self, offset, blocks, blocklight, skylight, data):
		# our true offset is based on the section's Y value
		offset += self.y * 16

		blocks[offset:offset + 16,:,:] = self.blocks
		blocklight[offset:offset + 16,:,:] = self.blocklight
		skylight[offset:offset + 16,:,:] = self.skylight
		data[offset:offset + 16,:,:] = self.data

# turns a 0-indexed layer into a Y offset
def get_layer_offset(layer):
	# with an offset of 64 and vanilla world height of 256
	# 0: 0 - 256
	# 1: 192 - 448
	# 2: 384 - 640
	# ...
	# n: n * (256 - 64) - ...
	
	return layer * (HEIGHT - OVERLAP)

# turns a 16x16x16 numpy array into an 8x16x16 compressed array
def compress_nibbles(a):
	dat = numpy.empty((8, 16, 16), dtype=numpy.uint8)
	dat[:,:,:] = (a[::2,:,:] & 0x0f) | ((a[1::2,:,:] & 0x0f) << 4)
	return dat

def convert_chunk(sources, num_layers, do_heightmap=False):
	# construct a basic template chunk, based on an arbitrary chunk
	template = list(sources.values())[0]
	converted = rs.NBT.new()

	# basic structure
	converted.root = rs.Tag.new0(rs.TAG_COMPOUND)
	level = rs.Tag.new0(rs.TAG_COMPOUND)
	converted.root['Level'] = level

	# copy over most info
	level['TerrainPopulated'] = template['Level']['TerrainPopulated']
	level['TileEntities'] = template['Level']['TileEntities']
	level['zPos'] = template['Level']['zPos']
	level['xPos'] = template['Level']['xPos']
	level['LastUpdate'] = template['Level']['LastUpdate']
	level['Entities'] = template['Level']['Entities']
	level['Biomes'] = template['Level']['Biomes']

	# continuous large arrays to render into
	final_height = get_layer_offset(num_layers - 1) + HEIGHT
	blocks = numpy.zeros((final_height, 16, 16), dtype=numpy.uint8)
	blocklight = numpy.zeros((final_height, 16, 16), dtype=numpy.uint8)
	skylight = numpy.ones((final_height, 16, 16), dtype=numpy.uint8) * 15
	data = numpy.zeros((final_height, 16, 16), dtype=numpy.uint8)

	# iterate through all chunks
	# trash_sections: how many sections to shave off on top and bottom
	# defaults to half of the overlap, so for 64 overlap: 2 sections
	trash_sections = (OVERLAP // 2) // 16
	for i, source in sources.items():
		offset = get_layer_offset(i)
		# iterate through all sections
		for sectiontag in source['Level']['Sections']:
			if i != 0 and sectiontag['Y'].value in range(trash_sections):
				continue # trash unless bottom layer
			if i != num_layers - 1 and sectiontag['Y'].value in range(16-trash_sections, 16):
				continue # trash unless top layer
			section = Section(sectiontag)
			# render this sucker
			section.render(offset, blocks, blocklight, skylight, data)
	# calculate the new HeightMap
	heightmap = numpy.zeros((16, 16), dtype=numpy.int32)
	if do_heightmap:
		for x in range(16):
			for z in range(16):
				y = final_height
				while y > 0 and blocks[y-1,x,z] == 0:
					y -= 1
				heightmap[x,z] = y

	# now write out the huge arrays
	tagheightmap = rs.Tag.new0(rs.TAG_INT_ARRAY)
	tagheightmap.value = list(heightmap.reshape((16 * 16,)))
	level['HeightMap'] = tagheightmap
	sections = rs.Tag.new0(rs.TAG_LIST)
	sections.list_set_type(rs.TAG_COMPOUND)
	level['Sections'] = sections
	for s in range(final_height // 16):
		secblocks = blocks[16 * s:16 * s + 16,:,:]
		secblocklight = blocklight[16 * s:16 * s + 16,:,:]
		secskylight = skylight[16 * s:16 * s + 16,:,:]
		secdata = data[16 * s:16 * s + 16,:,:]

		secblocklight = compress_nibbles(secblocklight)
		secskylight = compress_nibbles(secskylight)
		secdata = compress_nibbles(secdata)

		sec = rs.Tag.new0(rs.TAG_COMPOUND)
		sec['Y'] = rs.Tag.new0(rs.TAG_BYTE)
		sec['Y'].value = s

		tagblocks = rs.Tag.new0(rs.TAG_BYTE_ARRAY)
		tagblocks.value = secblocks.tostring()
		sec['Blocks'] = tagblocks

		tagblocklight = rs.Tag.new0(rs.TAG_BYTE_ARRAY)
		tagblocklight.value = secblocklight.tostring()
		sec['BlockLight'] = tagblocklight

		tagskylight = rs.Tag.new0(rs.TAG_BYTE_ARRAY)
		tagskylight.value = secskylight.tostring()
		sec['SkyLight'] = tagskylight

		tagdata = rs.Tag.new0(rs.TAG_BYTE_ARRAY)
		tagdata.value = secdata.tostring()
		sec['Data'] = tagdata

		sections.append(sec)
	
	return converted

if __name__ == "__main__":
	sys.exit(main())