SaveWizard/main/script.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-

from os import system, remove
from PyQt5.QtCore import Qt, QRegExp
from PyQt5.QtGui import QRegExpValidator
from PyQt5.QtWidgets import QMainWindow, QFileDialog

from .form import Ui_MainWindow
from util import *
from dataIO import dataIO
from second.script import SecondWindow


class MainWindow(QMainWindow, Ui_MainWindow):
    def __init__(self, parent=None):
        # Setup UI
        QMainWindow.__init__(self, parent, flags=Qt.Window)
        Ui_MainWindow.__init__(self)
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)

        self.file_path = ""
        self.old_file = ""

        # Checking DLC file
        if dataIO.is_valid_json("dlc.json") is False:
            self.owns = False
            show_message(QMessageBox.Warning, "Warning", "'dlc.json' not found, functionality has been limited")
        else:
            self.owns = {}
            self.dlc = dataIO.load_json("dlc.json")

        # Storing edits with his checkboxes and file-lines
        self.basic_edits = {
            self.ui.money_edit: [self.ui.money_dont_change, "money_account:"],
            self.ui.xp_edit: [self.ui.xp_dont_change, "experience_points:"],
            self.ui.loan_limit_edit: [self.ui.loan_limit_dont_change, "loan_limit:"],
        }
        self.skill_edits = {
            self.ui.long_distance_edit: [self.ui.long_distance_dont_change, "long_dist:"],
            self.ui.high_value_cargo_edit: [self.ui.high_value_cargo_dont_change, "heavy:"],
            self.ui.fragile_cargo_edit: [self.ui.fragile_cargo_dont_change, "fragile:"],
            self.ui.urgent_delivery_edit: [self.ui.urgent_delivery_dont_change, "urgent:"],
            self.ui.ecodriving_edit: [self.ui.ecodriving_dont_change, "mechanical:"],
        }

        # Setting up validators for edits
        basic_validator = QRegExpValidator(QRegExp("[0-9]{1,9}"))
        for key in self.basic_edits.keys():
            key.setValidator(basic_validator)
            key.textEdited.connect(self.text_edited)

        self.ui.adr_edit.textEdited.connect(self.text_edited)  # TODO: Validator for ADR
        skills_validator = QRegExpValidator(QRegExp("[0-6]{1,1}"))
        for key in self.skill_edits.keys():
            key.setValidator(skills_validator)
            key.textEdited.connect(self.text_edited)

        # Connecting buttons
        self.ui.path_button.clicked.connect(self.open_file_dialog)
        self.ui.apply.clicked.connect(self.apply_changes)
        self.ui.backup.clicked.connect(self.recover_backup)
        self.ui.second_window.clicked.connect(self.open_second_win)

    def text_edited(self):
        sender = self.sender()
        if sender in self.basic_edits:
            self.basic_edits[sender][0].setChecked(False)
        if sender == self.ui.adr_edit:
            self.ui.adr_dont_change.setChecked(False)
        if sender in self.skill_edits:
            self.skill_edits[sender][0].setChecked(False)

    @staticmethod
    def get_adr(value):
        bin_code = bin(int(value))[2:]
        bin_code = "0" * (6 - len(bin_code)) + bin_code
        r = []
        for i in bin_code:
            r.append(i)
        return r

    def get_adr_from_line(self):
        adr_list = list(self.ui.adr_edit.text())
        for i in adr_list:
            if (i == " ") or (i == ",") or (i == "."):
                adr_list.remove(i)
        return adr_list

    def clear_fields(self):
        self.file_path = ""
        self.old_file = ""
        set_lines([])

        if self.owns is not False:
            self.owns = {}

        for key, value in self.basic_edits.items():
            key.setText("")
            value[0].setChecked(True)

        self.ui.adr_edit.setText("")
        self.ui.adr_dont_change.setChecked(True)
        for key, value in self.skill_edits.items():
            key.setText("")
            value[0].setChecked(True)

        self.ui.apply.setEnabled(False)
        self.ui.backup.setEnabled(False)
        self.ui.second_window.setEnabled(False)

    def get_file_data(self, file):
        try:
            with open(file, "r") as f:
                self.old_file = f.read()
        except UnicodeDecodeError:
            try:
                system("SII_Decrypt.exe --on_file -i \"{}\"".format(file))
                with open(file, "r") as f:
                    self.old_file = f.read()
                show_message(QMessageBox.Information, "Success", "File successfully decrypted.")
            except UnicodeDecodeError:
                show_message(QMessageBox.Critical, "Error", "Error to decrypt and open file. Try again.")
                return
        set_lines(self.old_file.split("\n"))

        if self.owns is not False:
            self.owns["base"] = True
            for i in get_array_items(search_line("companies:")):
                for key, value in self.dlc.items():
                    if value in i:
                        self.owns[key] = True

        for key, value in self.basic_edits.items():
            key.setText(str(get_value(search_line(value[1]))))

        adr = self.get_adr(get_value(search_line("adr:")))
        adr_list = ""
        for i in range(6):
            adr_list += adr[i] + "," if i != 5 else adr[i]
        self.ui.adr_edit.setText(adr_list)

        for key, value in self.skill_edits.items():
            key.setText(str(get_value(search_line(value[1]))))

        self.ui.apply.setEnabled(True)
        self.ui.backup.setEnabled(True)
        self.ui.second_window.setEnabled(True)

    def open_file_dialog(self):
        file, _ = QFileDialog.getOpenFileName(parent=self,
                                              caption=self.tr("Choose your save file..."),
                                              filter=self.tr("game.sii"))
        self.clear_fields()
        if file != "":
            self.file_path = file
            self.get_file_data(self.file_path)
        else:
            return

    def apply_changes(self):
        if not self.ui.dont_change_all_inf.isChecked():
            for key, value in self.basic_edits.items():
                if value[0].isChecked() is False:
                    set_value(search_line(value[1]), key.text())
                    value[0].setChecked(True)
            if self.ui.adr_dont_change.isChecked() is False:
                adr_set = self.get_adr_from_line()
                if len(adr_set) < 6:
                    show_message(QMessageBox.Critical, "Error", "ADR can't have less than 6 elements.")
                elif len(adr_set) > 6:
                    show_message(QMessageBox.Critical, "Error", "ADR can't have more than 6 elements.")
                else:
                    adr_new = int("".join(adr_set), 2)
                    set_value(search_line("adr:"), str(adr_new))
            for key, value in self.skill_edits.items():
                if value[0].isChecked() is False:
                    set_value(search_line(value[1]), key.text())
                    value[0].setChecked(True)
        backup = self.file_path + ".swbak"
        with open(backup, "w") as f:
            f.write(self.old_file)
        with open(self.file_path, "w") as f:
            f.write("\n".join(get_lines()))
        show_message(QMessageBox.Information, "Success", "Changes successfully applied!")
        self.get_file_data(self.file_path)
        return

    def recover_backup(self):
        try:
            backup = self.file_path + ".swbak"
            f = open(backup, "r")
            with open(self.file_path, "w") as g:
                g.write(f.read())
            f.close()
            remove(backup)
            show_message(QMessageBox.Information, "Success", "Backup successfully recovered.")
            self.get_file_data(self.file_path)
        except IOError:
            show_message(QMessageBox.Critical, "Error", "Backup not found.")
            return

    def open_second_win(self):
        second_win = SecondWindow(self.owns, self)
        second_win.show()