giants-wdefs-importer-gui/main.py

import os
import struct
import sys
import traceback
from pathlib import Path
import shutil
import zlib
import requests
from PySide2.QtCore import QProcess
from PySide2.QtWidgets import QApplication, QMainWindow, QFileDialog, QMessageBox, QPushButton
from mainwindow import Ui_MainWindow


class MainWindow(QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)
        self.giants_path = locate_giants_folder()
        if not self.giants_path:
            raise Exception()
        self.wdefs_path = self.giants_path / "Bin" / "wdefs.bin"
        self.giantsmain_path = self.giants_path / "GiantsMain.exe"
        self.ui.lineEdit.textEdited.connect(self.on_textedit_changed)
        self.ui.pushButton.clicked.connect(self.on_button_clicked)

    def on_textedit_changed(self):
        if self.ui.lineEdit.text():
            self.ui.pushButton.setEnabled(True)
        else:
            self.ui.pushButton.setEnabled(False)

    def on_button_clicked(self):
        self.ui.pushButton.setEnabled(False)
        try:
            backup_file(self.wdefs_path)
            backup_file(self.giantsmain_path)
        except Exception:
            traceback.print_exc()
            QMessageBox.critical(None, "Wdefs", "Error while creating backups.")
            return

        text = self.ui.lineEdit.text()
        if not text:
            QMessageBox.critical(None, "Wdefs", "URL is empty")
            return

        try:
            req = requests.get(text)
            req.raise_for_status()
            start_giants_with_wdefs_bytes(req.content, self.wdefs_path, self.giantsmain_path, self.ui.pushButton)
        except requests.HTTPError:
            traceback.print_exc()
            QMessageBox.critical(None, "Wdefs", "Could not download wdefs at URL.")
        except Exception:
            try:
                revert_files(self.giantsmain_path)
                revert_files(self.wdefs_path)
                traceback.print_exc()
                QMessageBox.warning(None, "Wdefs", "There was a problem while applying mod. Your game has been restored to original state.")
                return
            except Exception:
                traceback.print_exc()
                QMessageBox.critical(None, "Wdefs", "There was a critical error while applying the mod. Your game couldn't be restored to original state. Please message the mods on Discord.")
                return


def start_giants_with_wdefs_bytes(wdefs_bytes: bytes, wdefs_path: Path, giantsmain_path: Path, btn: QPushButton):
    with open(wdefs_path, "wb") as fp:
        fp.write(wdefs_bytes)

    new_bytes = new_exe_bytes(wdefs_bytes, giantsmain_path)
    with open(giantsmain_path, "wb") as fp:
        fp.write(new_bytes)

    def giants_finished():
        revert_files(giantsmain_path)
        revert_files(wdefs_path)
        QMessageBox.information(None, "Wdefs", "Giants has been reverted to original state.")
        btn.setText("Play")
        btn.setEnabled(True)

    process = QProcess(app)
    process.finished.connect(giants_finished)
    process.setWorkingDirectory(str(giantsmain_path.parent))
    btn.setText("Game started")
    if sys.platform == "win32":
        process.start(str(giantsmain_path), ["-launcher"])
    else:
        process.start("/home/tasty/.local/share/lutris/runners/wine/lutris-6.0-x86_64/bin/wine", ["/home/tasty/Jeux/Giants Citizen Kabuto1.498/GiantsMain.exe", "-window", "-launcher"])


def revert_files(file_path: Path) -> None:
    # replace current with _original
    dst = str(file_path.parent / file_path.stem) + "_original" + file_path.suffix
    shutil.copy(dst, file_path)


def backup_file(file_path: Path) -> None:
    # copy and append _original to filename
    # if file already exists, skip
    dst = str(file_path.parent / file_path.stem) + "_original" + file_path.suffix
    print(dst)
    if os.path.exists(dst):
        return
    shutil.copy(file_path, dst)


def locate_giants_folder() -> Path:
    is_windows = sys.platform == "win32"
    if is_windows:
        possible_paths = [
            "C:\\Program Files (x86)\\GOG Galaxy\\Games\\Giants - Citizen Kabuto",
            "C:\\Program Files\\GOG Galaxy\\Games\\Giants - Citizen Kabuto",
            "C:\\Program Files (x86)\\GOG.com\\Games\\Giants - Citizen Kabuto",
            "C:\\Program Files\\GOG.com\\Games\\Giants - Citizen Kabuto",
            "C:\\Program Files (x86)\\Steam\\steamapps\\common\\Giants Citizen Kabuto",
            "C:\\Program Files\\Steam\\steamapps\\common\\Giants Citizen Kabuto",
        ]
        for path in possible_paths:
            if os.path.exists(path+"\\GiantsMain.exe"):
                return Path(path)

        import winreg
        try:
            path = winreg.QueryValue(winreg.HKEY_CURRENT_USER, "Software\\PlanetMoon\\Giants\\DestDir")
            if os.path.exists(path + "\\GiantsMain.exe"):
                return Path(path)
        except:
            pass
        return ask_giants_directory()
    else:
        return ask_giants_directory()


def ask_giants_directory() -> Path:
    QMessageBox.information(None, "Wdefs",
                            "Could not locate your Giants installation automatically.\nPlease browse to your Giants folder and select GiantsMain.exe")
    giantsmain_path, _filename = QFileDialog.getOpenFileName(None, caption="Wdefs", filter="GiantsMain.exe")
    if giantsmain_path:
        return Path(giantsmain_path).parent
    return None


def crc(inp_bytes: bytes) -> int:
    prev = zlib.crc32(inp_bytes, 0)
    return prev & 0xffffffff


def new_exe_bytes(wdefs_bytes: bytes, giantsmain_path: Path) -> bytes:
    CHECKSUM_1_502_0_1 = b"\x78\x05\x44\xb6"  # should be at offset 0x152bff in giantsmain.exe

    with open(giantsmain_path, "rb") as fp:
        content = fp.read()

    index = content.find(CHECKSUM_1_502_0_1)
    if index < 0:
        raise Exception("Could not find wdefs checksum in embedded GiantsMain.exe. Report this to Amazed#0001")
    else:
        content = bytearray(content)
        wdefs_crc = struct.pack("<L", crc(wdefs_bytes))
        content[index:index+4] = wdefs_crc
        return content


if __name__ == "__main__":
    app = QApplication(sys.argv)
    app.setQuitOnLastWindowClosed(True)
    try:
        window = MainWindow()
        window.show()
        sys.exit(app.exec_())
    except SystemExit:
        sys.exit(0)
    except:
        traceback.print_exc()
        sys.exit(1)