path: root/ModSlots/ModSlots.py

#! /usr/bin/env python3

import sys
sys.dont_write_bytecode = True
import os.path
import hashlib
from enum import Enum
from zipfile import ZipFile
#from unrardll import extract, names
#from py7zr import SevenZipFile
from ArmorIDS import ArmorIDS

# @TODO:
# - Detach files (copy instead of symlink, MHWNewCamera.json)
# - Test example to normal and back
# - More automatic override for the simplest case
# - Support invisible weapons (default no-include)

USAGE = f'Usage: {sys.argv[0]} <Write/Verify/Delete> <Slots.txt> <Path_to_Mods> <Path_to_Game>'

if len(sys.argv) < 4:
    print(USAGE)
    sys.exit(1)

Mode = sys.argv[1].lower()
Slots = sys.argv[2]
Mods_Directory = sys.argv[3]
Game_Directory = sys.argv[4]
Extract_Directory = f'{Mods_Directory}/extract'
Hash_Path = f'{Game_Directory}/mod_slots_hash.txt'
List_Path = f'{Game_Directory}/mod_slots.txt'

if not os.path.isfile(Slots):
    print(f'{Slots} not a file.')
    sys.exit(1)

if not os.path.isdir(Mods_Directory):
    print(f'Directory {Mods_Directory} doesn\'t exist.')
    sys.exit(1)

if not os.path.isdir(Extract_Directory):
    try:
        os.mkdir(Extract_Directory)
    except:
        print(f'Failed to create directory for extracted mods at {Extract_Directory}.')
        sys.exit(1)

if not os.path.isdir(Game_Directory):
    print(f'Directory {Game_Directory} doesn\'t exist.')
    sys.exit(1)

Error_Messages = []
Error_Index = 1
def Error(msg):
    global Error_Messages
    global Error_Index
    msg = f'{Error_Index}: ' + msg
    Error_Messages.append(msg)
    Error_Index += 1
    print(msg + '\n^^^^^^^^^^^^^^^^^^^^^^^')

ArmorIDS_Inverted = { v: k for k, v in ArmorIDS.items() }

class ArchiveType(Enum):
    ZIP = 0
    RAR = 1
    SEVEN_ZIP = 2

class Mod():
    def __init__(self, path):
        m = hashlib.sha256()
        m.update(open(path, 'rb').read())
        self.hash = m.hexdigest()
        ext = path.split('.')[-1].lower()
        if ext == 'zip':
            self.type = ArchiveType.ZIP
            self.obj = ZipFile(path, 'r')
            self.files = []
            for file in self.obj.infolist():
                if file.filename[-1] == '/':
                    continue
                self.files.append(file.filename)
        elif ext == 'rar':
            self.type = ArchiveType.RAR
            pass
        elif ext == '7z':
            self.type = ArchiveType.SEVEN_ZIP
            pass
        self.toplevels = []
        self.roots = []
        self.ignored = []
        self.overrides = {}
        self.map = {}

    def substitute_override_name(self, name):
        if name[0] == '[':
            name = name[1:-1]
            if name in ArmorIDS_Inverted.keys():
                name = ArmorIDS_Inverted[name][2:]
            else:
                return None
        return name

    def get_id_from_subpath(self, sub):
        for armor_id in ArmorIDS.keys():
            armor_id = armor_id[2:]
            search = sub.find(armor_id)
            if search >= 0:
                return sub[search:search + 8]
        return None

    def find_overrides_for_file(self, file):
        for key in self.overrides.keys():
            search = file.find(key)
            if search >= 0:
                return self.overrides[key]
        return None

    def file_sort_key(self, x):
        for i in range(0, len(self.roots)):
            if x.startswith(self.roots[i]):
                return i
        return 0

    def do_map(self, All_Files):
        for file in sorted(self.files, key=self.file_sort_key):
            ignore_this = False
            for ig in self.ignored:
                if file.startswith(ig):
                    ignore_this = True
                    break
            if ignore_this:
                continue
            path_in_archive = file
            # If file is within a defined root, truncate it's path up to nativePC/.
            for root in self.roots:
                if file.startswith(root):
                    nativePC = file.find('nativePC')
                    if nativePC >= 0:
                        file = file[nativePC:]
                        break
            targets = []
            if file in self.overrides.keys():
                targets.extend(self.overrides[file])
            elif not (file.startswith('nativePC') or file in self.toplevels):
                # If file isn't within nativePC/ and is not explicitly included, ignore it.
                continue
            else:
                overrides = self.find_overrides_for_file(file)
                if not overrides:
                    targets.append(file) # Include file as-is.
                else:
                    # Substitute armor names, like '[Anjanath]', with their IDs.
                    for override in overrides:
                        target_id = self.get_id_from_subpath(file)
                        if not target_id:
                            Error(f'Invalid target: {file}.')
                            continue
                        mapped_name = self.substitute_override_name(override)
                        if not mapped_name:
                            Error(f'Invalid substitute: {override}.')
                            continue
                        targets.append(file.replace(target_id, mapped_name))
            print(path_in_archive)
            unique_targets = []
            for target in targets:
                # @TODO: Allow "important" mapping instead of relying on order?
                if target in All_Files:
                    print(f' (Overwritten)')
                else:
                    unique_targets.append(target)
                    print(' -> ' + target)
            if len(unique_targets) > 0:
                All_Files.extend(unique_targets)
                self.map[path_in_archive] = unique_targets

    def add_to_hash(self, m):
        output = f'{os.path.join(Extract_Directory, self.hash)}'
        if not os.path.isdir(output):
            os.mkdir(output)
            if self.type == ArchiveType.ZIP:
                self.obj.extractall(path=output)
        for key in self.map:
            m.update(open(os.path.join(output, key), 'rb').read())

    def write(self, old_list):
        output = f'{os.path.join(Extract_Directory, self.hash)}'
        with open(List_Path, 'a+') as l:
            for key in self.map:
                targets = self.map[key]
                for target in targets:
                    if target in old_list:
                        old_list.remove(target)
                    l.write(target + '\n')
                    inc = Game_Directory
                    for s in target.split('/')[:-1]:
                        inc = os.path.join(inc, s)
                        if not os.path.isdir(inc):
                            os.mkdir(inc)
                    target = os.path.join(Game_Directory, target)
                    if os.path.islink(target):
                        os.remove(target)
                    elif os.path.isfile(target) or os.path.isdir(target):
                        Error(f'Not deleting non-symlink file: {target}.')
                        continue
                    os.symlink(f'{os.path.join(output, key)}', target)

    def delete(self):
        for key in self.map:
            targets = self.map[key]
            for target in targets:
                target = os.path.join(Game_Directory, target)
                if os.path.islink(target):
                    os.remove(target)
                elif os.path.isfile(target) or os.path.isdir(target):
                    Error(f'Not deleting non-symlink file: {target}.')

def is_armor_override(line):
    sp = line.split('/')
    if len(sp) >= 3 and sp[2] == 'slg':
        return True
    if sp[-1] in ['helm', 'body', 'arm', 'wst', 'leg']:
        return True
    return False

Mods = []
Current_Mod = None
Current_Override = None
with open(Slots, 'r') as f:
    for line in f.read().splitlines():
        if len(line) == 0 or line[0] == '#':
            continue
        if line[0] == ';':
            if Current_Mod:
                Mods.append(Current_Mod)
                Current_Mod = None
            path = os.path.join(Mods_Directory, line[1:])
            if os.path.isfile(path):
                Current_Mod = Mod(path)
            else:
                Error(f'File not found: {path}')
            continue
        cmd = ''
        if line[0] == '*':
            cmd = 'root'
        elif line[0] == '<':
            cmd = 'ignore'
        elif line[0] == ':':
            cmd = 'override'
        elif line[0] == '>':
            cmd = 'value'
        if cmd:
            line = line[1:]
            if line.startswith('pl/'):
                line = 'nativePC/' + line
            if line.startswith('wp/'):
                line = 'nativePC/' + line
        if cmd == 'root' and Current_Mod:
            Current_Mod.roots.append(line)
        elif cmd == 'ignore' and Current_Mod:
            Current_Mod.ignored.append(line)
        elif cmd == 'override' and Current_Mod:
            if not (is_armor_override(line) or line in Current_Mod.files):
                Error(f'File not found: {line}')
                Current_Mod = None
                continue
            Current_Override = line
            Current_Mod.overrides[Current_Override] = []
        elif cmd == 'value' and Current_Mod:
            if Current_Override:
                Current_Mod.overrides[Current_Override].append(line)
            else:
                Current_Mod.toplevels.append(line)
    if Current_Mod:
        Mods.append(Current_Mod)
        Current_Mod = None

All_Files = []
M = hashlib.sha256()
for mod in Mods:
    mod.do_map(All_Files)
    mod.add_to_hash(M)

with open(List_Path, 'a+') as l:
    l.seek(0)
    old_list = l.read().splitlines()
    l.truncate(0)

if Mode == 'write':
    for mod in Mods:
        mod.write(old_list)
    for old in old_list:
        old_target = os.path.join(Game_Directory, old)
        if os.path.islink(old_target):
            print(f'Removing old link: {old_target}')
            os.remove(old_target)
    with open(Hash_Path, 'w+') as f:
        f.write(M.hexdigest())
if Mode == 'verify':
    if os.path.isfile(Hash_Path):
        with open(Hash_Path, 'r') as f:
            if f.read() == M.hexdigest():
                print('Verification passed')
            else:
                print('Verification failed')
elif Mode == 'delete':
    for mod in Mods:
        mod.delete()

if len(Error_Messages) > 0:
    print('-----------------------')
    for err in Error_Messages:
        print(err)