serialize-ubergraph.py

#!/usr/bin/env python3
# vim: set expandtab tabstop=4 shiftwidth=4:

#uassetapi_dir = '/home/pez/Programs/games/borderlands_3/uassetgui/v1.0.0.0-alpha.6'
#uassetapi_dir = '/home/pez/git/b2patching/UAssetAPI/UAssetAPI/bin/Debug'
uassetapi_dir = '/home/pez/git/b2patching/UAssetAPI/UAssetAPI/bin/Release'

import os
import clr
import sys
import argparse
import pythonnet
clr.AddReference(os.path.join(uassetapi_dir, 'UAssetAPI'))
import UAssetAPI

def get_serializations(filename):
    ass = UAssetAPI.UAsset(
            path=filename,
            engineVersion=UAssetAPI.UnrealTypes.UE4Version.VER_UE4_20,
            )

    UAssetAPI.Kismet.KismetSerializer.asset = ass
    for idx, export in enumerate(ass.Exports):
        if hasattr(export, 'ScriptBytecode'):
            if export.ScriptBytecode:
                # Attempt serialization (return is an object, but stringifies nicely)
                serialized = UAssetAPI.Kismet.KismetSerializer.SerializeScript(export.ScriptBytecode)
                yield (idx, export.ObjectName, serialized)

def main():

    parser = argparse.ArgumentParser(
            description='Serialize Ubergraph Bytecode using UAssetAPI',
            )

    parser.add_argument('-r', '--runtime',
            action='store_true',
            help='Show .NET runtime being used',
            )

    parser.add_argument('filename',
            type=str,
            nargs=1,
            help='Filename to process',
            )

    args = parser.parse_args()
    args.filename = args.filename[0]

    if args.runtime:
        print(pythonnet.get_runtime_info())

    obj_exts = {'uasset', 'umap'}
    _, filename_alone = os.path.split(args.filename)
    if '.' in filename_alone:
        filename_base, ext = args.filename.rsplit('.', 1)
    else:
        filename_base = args.filename
        ext = ''
    if ext not in obj_exts:
        for ext in obj_exts:
            if os.path.exists(f'{filename_base}.{ext}'):
                args.filename = f'{filename_base}.{ext}'
                break
    if not os.path.exists(args.filename):
        raise RuntimeError(f'Not found: {args.filename}')

    for index, name, serialization in get_serializations(args.filename):
        to_filename = f'{filename_base}-ubergraph-{index+1:03d}-{name}.json'
        with open(to_filename, 'w') as odf:
            odf.write(str(serialization))
        print(f'Wrote to: {to_filename}')

if __name__ == '__main__':
    main()