gsm_mapping_generator.py

# gsm_mapping_generator.py
import xml.etree.ElementTree as ET
import xml.dom.minidom
import os
from special_path_resolver import resolve_special_path
import re
import yaml

def iterate_subfolders(folderpath):
    folders = []
    while True:
        folderpath, subfolder = os.path.split(folderpath)
        if subfolder:
            folders.append(subfolder)
        else:
            if folderpath:
                folders.append(folderpath)
            break
    folders.reverse()
    return folders


def create_gsm_info_xml_mapping(yaml_contents: dict, output_folder_path: str, registry_yaml_contents = None):

    # Create the root element (e.g. specific version)
    root = ET.Element('GameSaveManager_EntryData', version="3.1.442.0 (Database: 9.2016-12-24_01)", attrib_hashes="true")

    # Add GameName element
    ET.SubElement(root, 'GameName', IsOfficial="true").text = yaml_contents['name']
    
    date_string = yaml_contents['backups'][0]['when']
        
    inner_backup_folder_name = yaml_contents['backups'][0]['name']

    # Add Warning element
    ET.SubElement(root, 'Warning')
    files = dict(yaml_contents['backups'][0])['files']
    
    #No dupes in set (then turn to list to sort)
    folders = list(set([os.path.dirname(file) for file in files]))
    folders.sort()

    # Iterate over backups
    for folder in folders:

        # Add Directory element for each backup
        directory = ET.SubElement(root, 'Directory')
        folderpath = folder.replace("/","\\")
        subfolders = iterate_subfolders(folderpath)
        current_folder = subfolders[0]
        candidate_special_paths = []
        rest_of_subfolders = None
        for subfolder in subfolders:
            new_folder = os.path.join(current_folder, subfolder).replace("/", "\\")
            current_folder = new_folder.replace("/","\\")
            possible_special_folder = resolve_special_path(current_folder, from_values=True)
            #if possible_special_folder is not None:
            if isinstance(possible_special_folder, str):
                candidate_special_paths.append(possible_special_folder)
                rest_of_subfolders = folderpath.replace(current_folder, "")
                if rest_of_subfolders.startswith("\\"):
                    rest_of_subfolders = rest_of_subfolders[1:]

        probable_special_path = ""
        if len(candidate_special_paths) == 1:
            probable_special_path = candidate_special_paths[0]
        else:
            for i in range(1, len((candidate_special_paths))):
                if resolve_special_path(candidate_special_paths[i]) != None and resolve_special_path(candidate_special_paths[i-1]) == None: 
                    probable_special_path = candidate_special_paths[i]
                elif resolve_special_path(candidate_special_paths[i]) == None and resolve_special_path(candidate_special_paths[i-1]) != None:
                    probable_special_path = candidate_special_paths[i-1]
                elif resolve_special_path(candidate_special_paths[i]) != None and resolve_special_path(candidate_special_paths[i-1]) != None:
                    if len(resolve_special_path(candidate_special_paths[i-1])) < len(resolve_special_path(candidate_special_paths[i])):
                        probable_special_path = candidate_special_paths[i]
                    else:
                        probable_special_path = candidate_special_paths[i-1]

        # Add SpecialPath, Path, and RegData elements
        if probable_special_path == None or probable_special_path == "":
            ET.SubElement(directory, 'SpecialPath')
            special_path_element = None
        else:               
            special_path_element = ET.SubElement(directory, 'SpecialPath').text = "%" + probable_special_path + "%"
        
        uplay_id = None
        steam_id = None

        if rest_of_subfolders != None:
            if probable_special_path == "UPLAY" or probable_special_path == "UPLAY_2":
                path_element = ET.SubElement(directory, 'Path').text = rest_of_subfolders.rsplit("\\")[1]
                uplay_id = rest_of_subfolders.rsplit("\\")[0]
            elif probable_special_path == "STEAM_CLOUD":
                path_element = ET.SubElement(directory, 'Path').text = rest_of_subfolders.rsplit("\\")[1]
                steam_id = rest_of_subfolders.rsplit("\\")[0]
            else:
                path_element = ET.SubElement(directory, 'Path').text = rest_of_subfolders.rsplit("\\", 1)[0] # backup['name']
        else:
            path_element = ET.SubElement(directory, 'Path').text = folder
            
        reg_data = ET.SubElement(directory, 'RegData')
        ET.SubElement(reg_data, 'Hive')
        ET.SubElement(reg_data, 'Path')
        ET.SubElement(reg_data, 'Value')

        file_list = ET.SubElement(directory, 'FileList')
        file_paths = files.keys()
        total_folder_size = 0

        if probable_special_path == None or special_path_element in {"%%","",None}:
            fullprefix_path = path_element
        else:
            if uplay_id != None:
                fullprefix_path = os.path.join(resolve_special_path(special_path_element), uplay_id, path_element)
            elif steam_id != None:
                fullprefix_path = os.path.join(resolve_special_path(special_path_element), steam_id, path_element)
            else:
                fullprefix_path = os.path.join(resolve_special_path(special_path_element), path_element)
            
        fullprefix_path = fullprefix_path.replace("/","\\")
        for file_path in file_paths:
            folder_of_file = os.path.dirname(file_path).replace("/","\\")
            if folder_of_file == folderpath:
                file_path_reformatted = file_path.replace("/","\\")
                ET.SubElement(file_list, 'File', checksum=files[file_path]['hash']).text = file_path_reformatted.replace(fullprefix_path, "")[1::] #os.path.basename(file_path)
                total_folder_size += files[file_path]['size']
                
        # Add Size element
        ET.SubElement(directory, 'Size').text = str(total_folder_size)

        # Add UserInfo element
        user_info = ET.SubElement(directory, 'UserInfo')
        if steam_id != None:
            ET.SubElement(user_info, 'SteamID').text = steam_id
        else:
            ET.SubElement(user_info, 'SteamID')
        ET.SubElement(user_info, 'SteamUser')
        ET.SubElement(user_info, 'FlashID')
        if uplay_id != None:
            ET.SubElement(user_info, 'Uplay').text = uplay_id
        else:
            ET.SubElement(user_info, 'Uplay')
    
    game_name = (yaml_contents['name'].replace(":","")).replace('â„¢', '™')
    specific_output_folder = os.path.join(output_folder_path, f"{game_name}" + "_" + date_string[0:19].replace("T","_").replace(":","."))
    #specific_output_folder = specific_output_folder + "_" + date_string[0:19].replace("T","_").replace(":",".")

    
    # Registry data processing
    #if yaml_contents['backups'][0]['registry']['hash'] != None and registry_yaml_contents != None:
    if yaml_contents['backups'][0]['registry']['hash'] and registry_yaml_contents:
        # Parse the YAML contents
        registry_data = yaml.safe_load(registry_yaml_contents)

        # Extract the data
        root_keys = list(registry_data.keys())
          

        # Format the data into .reg file format
        reg_file_contents = f"""Windows Registry Editor Version 5.00\n\n"""
        for root_key in root_keys:
            sub_keys = list(registry_data[root_key].keys())#[0]
            for sub_key in sub_keys:
                value_names = list(registry_data[root_key][sub_key].keys())#[0]
                reg_file_contents += f"[{root_key}\\{sub_key}]\n"
                for value_name in value_names:
                    values_type = list(registry_data[root_key][sub_key][value_name].keys())#[0]
                    for value_type in values_type:
                        #value = list(registry_data[root_key][sub_key][value_name][value_type])#[0]
                        value = registry_data[root_key][sub_key][value_name][value_type]#[0]
                        #for value in values:
                        if value_type == "binary":
                            value_type = "hex"
                        if isinstance(value, list):
                            value = ','.join([format(i, '02x') for i in value])
                        elif isinstance(value, int):
                            value = format(value, '08x') #hex(value)
                        if value_type == "sz":
                            reg_file_contents += f'"{value_name}"="{value}"\n'
                        else:
                            reg_file_contents += f'"{value_name}"={value_type}:{value}\n'
                reg_file_contents += "\n"
            reg_file_contents += "\n"
        
        while reg_file_contents.endswith("\n"):
            reg_file_contents = reg_file_contents[:-1]

        # Write the .reg file
        os.makedirs(specific_output_folder, exist_ok=True)
        #with open(os.path.join(output_folder_path, f"{game_name}", f"{game_name}"+".reg"), 'w') as f:
        with open(os.path.join(specific_output_folder, f"{game_name}"+".reg"), 'w') as f:
            f.write(reg_file_contents)
    
    # Convert the XML structure to a string
    xml_string = ET.tostring(root, encoding='utf-8').decode('utf-8')
    
    # Parse the XML string with minidom and pretty-print it (hiding the 'xml version tag')
    dom = xml.dom.minidom.parseString(xml_string)
    pretty_xml_string = dom.childNodes[0].toprettyxml()

    # Replace all self-closing tags with open and close tags
    pretty_xml_string = re.sub(r'<(.+?)/>', r'<\1></\1>', pretty_xml_string)

    # Remove leading whitespace from each line
    pretty_xml_string = '\n'.join(line.lstrip() for line in pretty_xml_string.split('\n'))



    # Write the pretty-printed XML string to a file
    os.makedirs(specific_output_folder, exist_ok=True)
    
    #Check if output folder doesn't have any subfolders with mapping.yaml files inside or if any subfolders' name begins with "backup-"
        
    for root, dirs, files in os.walk(specific_output_folder):
        for file in files:            
            #Compare just file name, not the full path
            if os.path.basename(file) == "mapping.yaml":
                raise ValueError(f"ERROR: Output folder {specific_output_folder} contains subfolders with 'mapping.yaml' files. These may be Ludusavi backups, that may get overwritten. Please remove them before running the script.")
        for dir in dirs:
                if dir.startswith("backup-"):
                    raise ValueError(f"ERROR: Output folder {specific_output_folder} contains subfolders with names starting with 'backup-'. These may be Ludusavi backups, that may get overwritten. Please remove them before running the script.")
    
    
    with open(os.path.join(specific_output_folder, "GSM_INFO.xml"), 'w') as f:
        f.write(pretty_xml_string)
        
    return pretty_xml_string, specific_output_folder, date_string, inner_backup_folder_name