Wiizard/wiizard/source.py

########################################################################
#  Wiizard - A Wii games manager
#  Copyright (C) 2023  CYBERDEViL
#
#  This file is part of Wiizard.
#
#  Wiizard is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  Wiizard is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
########################################################################


import os

import pywwt

from wiizard.fileInfo import LocalGameFileInfo, FILE_TYPE_ISO, FILE_TYPE_WBFS
from wiizard.thread import AbstractThread, THREAD_FLAG_IS_CANCELLABLE
import wiizard.const as Const


class PartitionRepairThread(AbstractThread):
  def __init__(self, path, repairMode=pywwt.REPAIR_ALL, testMode=True):
    AbstractThread.__init__(self, flags=0)
    self.__path       = path
    self.__testMode   = testMode
    self.__repairMode = repairMode
    self.result       = None
    self.error        = ""

  def run(self):
    result = None

    try:
      result = pywwt.repair_wbfs_partition(self.__path,
                                           repairMode=self.__repairMode,
                                           testMode=self.__testMode)
    except pywwt.error as err:
      print("Repair failed:", err)
      self.error = str(err)

    self.result = result
    self.completed.emit()


class PartitionScanThread(AbstractThread):
  def __init__(self, path):
    AbstractThread.__init__(self, flags=0)
    self.__path = path
    self.result = None
    self.error  = ""

  def run(self):
    self.result = {}

    try:
      self.result.update(pywwt.check_wbfs_partition(self.__path))
    except pywwt.error as err:
      self.error += f"Checking WBFS partition failed: {err}"
      self.completed.emit()
      return

    try:
      self.result.update(pywwt.get_wbfs_partition_info(self.__path, discs=True))
    except pywwt.error as err:
      self.error += f"Failed to get WBFS partition info: {err}"

    self.completed.emit()


class InvalidFile:
  @staticmethod
  def isValid():
    return False


""" Read the ID6 and title from .iso and .wbfs files
"""
def getLocalGameFileInfo(path, ext, sourceModel):
  offset   = 0
  fileType = FILE_TYPE_ISO

  if ext == ".wbfs":
    fileType = FILE_TYPE_WBFS
    offset   = 512

  with open(path, 'rb') as fd:
    fd.seek(offset)
    try:
      id6 = fd.read(6).decode("ascii")
    except UnicodeDecodeError:
      return InvalidFile
    fd.seek(offset + 32)
    try:
      title = fd.read(64).decode("utf-8").rstrip("\x00")
    except UnicodeDecodeError:
      return InvalidFile
    fd.seek(0, os.SEEK_END)
    size = round(fd.tell() / Const.MiB, 2)
    return LocalGameFileInfo(id6, title, os.path.dirname(path),
                             os.path.basename(path), fileType, size,
                             sourceModel)


""" Recursive find all files with given extensions on local disk (non-wbfs)
"""
class FileFinder(AbstractThread):
  def __init__(self, path, sourceModel, extensions=Const.EXTS_WII, maxDepth=4):
    AbstractThread.__init__(self, flags=THREAD_FLAG_IS_CANCELLABLE)
    self.__path        = path
    self.__exts        = extensions
    self.__maxDepth    = maxDepth
    self.__sourceModel = sourceModel
    self.result        = None

  def run(self):
    self.result = self.findFiles(self.__path, self.__exts, maxDepth=self.__maxDepth)
    self.completed.emit()

  def findFiles(self, path, extensions=Const.EXTS_WII, depth=0, maxDepth=4):
    files = []

    try:
      for file_ in os.listdir(path):
        if self.cancelled:
          break

        realPath = os.path.join(path, file_)
        if os.path.isdir(realPath):
          if depth < maxDepth:
            files += self.findFiles(realPath, depth=depth+1, maxDepth=maxDepth)
          continue

        for ext in extensions:
          if file_.endswith(ext):
            fileInfo = getLocalGameFileInfo(realPath, ext, self.__sourceModel)
            if fileInfo.isValid():
              files.append(fileInfo)
            break
    except PermissionError:
      pass # TODO

    return files