export_record_metadata.py

# coding=utf-8

import io
import json
import os
import pathlib
import struct

from collections import defaultdict

OUTPUT_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'reports')

GAME_DATA_PATH = r'E:\Program Files (x86)\Steam\steamapps\common\Skyrim Special Edition\Data'


def generate_blocks(file):
	blocks = []
	with open(file, mode='rb') as f:
		file_extension = os.path.splitext(os.path.basename(file))[1]
		if file_extension not in ['.esp', '.esm']:
			raise ValueError('File extension did not match expected value: .esp, or .esm', file_extension)

		file_header = struct.unpack_from('4sii8xi', f.read(24))

		signature = file_header[0]
		data_size = file_header[1]
		# form_version = data[3]

		if signature != b'TES4':
			raise ValueError('File magic did not match expected value: b\'TES4\'', signature)
		else:
			# skip file header
			f.seek(data_size + 24, os.SEEK_SET)

		signature = f.read(4)
		if signature != b'GRUP':
			raise ValueError('Signature did not match expected value: b\'GRUP\'', signature)
		else:
			f.seek(-4, 1)

			while True:
				if f.tell() >= os.path.getsize(file):
					break

				data = struct.unpack_from('4si4s', f.read(12))
				block_size = data[1]

				f.seek(-12, 1)

				block = f.read(block_size)
				group = io.BytesIO(block)
				blocks.append(group)

	return blocks


def generate_json(blocks):
	records = []
	for block in blocks:
		block_header = struct.unpack_from('4si4s', block.read(12))
		block_size = block_header[1]
		block_signature = block_header[2]
		if block_signature in [b'CELL', b'WRLD', b'DIAL']:
			continue

		block.seek(12, 1)

		record_info = defaultdict(list)
		while True:
			if block.tell() >= block_size:
				break

			record_header = struct.unpack('4si4xi4xh2s', block.read(24))
			block.seek(-4, 1)

			record_signature = record_header[0].decode('utf-8')
			record_size = record_header[1]
			record_form_id = '{:08x}'.format(record_header[2])
			record_form_version = record_header[3]

			record_info[record_signature].append({'Form ID': record_form_id.upper(), 'Form Version': record_form_version, 'Size': record_size})

			block.seek(record_size + 4, 1)

		records.append(dict(record_info))

	return records


def main():
	files = [str(x) for x in pathlib.Path(GAME_DATA_PATH).glob('**/*.esm')]
	files += [str(x) for x in pathlib.Path(GAME_DATA_PATH).glob('**/*.esp')]

	for file in files:
		print('Exporting record metadata from file: ' + file)
		blocks = generate_blocks(file)
		json_object = generate_json(blocks)

		output_file_name = os.path.splitext(os.path.basename(file))[0]
		os.makedirs(OUTPUT_PATH, exist_ok=True)
		with open(os.path.join(OUTPUT_PATH, output_file_name + '.json'), 'w') as f:
			json.dump(json_object, f, indent=2)


if __name__ == '__main__':
	main()