ModSlots better extract logic, add shuffle

Signed-off-by: Andrew Opalach <andrew@akon.city>

1 files changed, 146 insertions, 65 deletions

diff --git a/ModSlots/ModSlots.py b/ModSlots/ModSlots.py
index 5a4f216..6460f4c 100755
--- a/ModSlots/ModSlots.py
+++ b/ModSlots/ModSlots.py
@@ -6,12 +6,15 @@ import os
 import os.path
 import shutil
 import re
+import struct
+import time
+import random
 from enum import Enum
 from hashlib import sha256
 from zipfile import ZipFile
 
 try:
-    from unrardll import extract, names
+    import unrardll
     RAR_SUPPORT = True
 except ImportError:
     RAR_SUPPORT = False
@@ -24,6 +27,9 @@ except ImportError:
 
 from ArmorIDs import ArmorIDs
 
+def flatten(xss):
+    return [x for xs in xss for x in xs]
+
 NATIVEPC_SUBDIRS = ('ec', 'hm', 'npc', 'pg', 'posteffect', 'sound', 'ui', 'village', 'accessory', 'em', 'light', 'otomo', 'photo', 'quest', 'system', 'unit_resource', 'common', 'event', 'ngword_list', 'pc', 'pl', 'shortcut', 'vfx', 'wp', 'plugins')
 
 USAGE = f'Usage: {sys.argv[0]} <Write/Verify/Delete> <Slots.txt> <Path_to_Mods> <Path_to_Game>'
@@ -37,7 +43,6 @@ 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):
@@ -59,8 +64,8 @@ if not os.path.isdir(Game_Directory):
     print(f'Directory {Game_Directory} doesn\'t exist.')
     sys.exit(1)
 
-if Mode == 'verify' and not os.path.isfile(Hash_Path):
-    print(f'Verification requested but previous hash not found at {Hash_Path}.')
+if Mode == 'verify' and not os.path.isfile(List_Path):
+    print(f'Verification requested but previous hash not found at {List_Path}.')
     sys.exit(1)
 
 Error_Messages = []
@@ -68,11 +73,20 @@ Error_Index = 1
 def Error(msg):
     global Error_Messages
     global Error_Index
-    msg = f'{Error_Index}: {msg}'
+    msg = f'E(#{Error_Index}): {msg}'
     Error_Messages.append(msg)
     Error_Index += 1
     print(msg + '\n' + '^'*len(msg))
 
+Warning_Index = 1
+def Warn(msg):
+    global Error_Messages
+    global Warning_Index
+    msg = f'W(#{Warning_Index}): {msg}'
+    Error_Messages.append(msg)
+    Warning_Index += 1
+    print(msg + '\n' + '^'*len(msg))
+
 class ArchiveType(Enum):
     ZIP = 0
     RAR = 1
@@ -80,6 +94,7 @@ class ArchiveType(Enum):
 
 class Mod():
     def __init__(self, path):
+        self.path = path
         self.name = path.split('/')[-1]
         m = sha256()
         m.update(open(path, 'rb').read())
@@ -94,17 +109,32 @@ class Mod():
                     continue
                 self.files.append(file.filename)
         elif ext == 'rar':
-            if not SEVEN_ZIP_SUPPORT:
-                raise Exception("Use of .rar archive without required package (unrardll).")
+            if not RAR_SUPPORT:
+                raise Exception('Use of .rar archive without required package (unrardll).')
             self.type = ArchiveType.RAR
+            for header in list(unrardll.headers(path)):
+                if not header['is_dir']:
+                    self.files.append(header['filename'])
         elif ext == '7z':
             if not SEVEN_ZIP_SUPPORT:
-                raise Exception("Use of .7z archive without required package (py7zr).")
+                raise Exception('Use of .7z archive without required package (py7zr).')
             self.type = ArchiveType.SEVEN_ZIP
             self.obj = SevenZipFile(path, 'r')
             for file in self.obj.list():
                 if not file.is_directory:
                     self.files.append(file.filename)
+        # Filter paths to fix automatic processing.
+        self.raw_files = self.files.copy()
+        for i, file in enumerate(self.files):
+            npc = file.lower().find('nativepc')
+            if npc >= 0:
+                cut = file[npc:npc + len('nativePC')]
+                if cut != 'nativePC':
+                    print(f'Correcting nativePC capitalization in {file}.')
+                    self.files[i] = file.replace(cut, 'nativePC')
+        self.raw_map = {}
+        for i in range(0, len(self.files)):
+            self.raw_map[self.files[i]] = self.raw_files[i]
         self.roots = []
         self.ignored = []
         self.overrides = {}
@@ -136,9 +166,9 @@ class Mod():
     def maybe_strip_root(self, file):
         for root in self.roots:
             if file.startswith(root):
-                nativePC = file.find('nativePC')
-                if nativePC >= 0:
-                    return file[nativePC:], root
+                npc = file.find('nativePC')
+                if npc >= 0:
+                    return file[npc:], root
         return file, None
 
     def do_map(self, All_Files, Global_Overrides):
@@ -163,7 +193,8 @@ class Mod():
                     if len(target.split('/')[-1].split('.')) == 1:
                         print(f'Warning: Possible non-direct path as a direct override: {target}')
                     targets.append(target)
-            elif not has_global_override: # Could switch to else.
+            # Add file as-is unless it could override an explicit target of a different file.
+            elif base not in flatten(self.overrides.values()) and not has_global_override:
                 targets.append(base)
             # Substitute armor names, like '[Anja]', with their IDs.
             for i, target in enumerate(targets):
@@ -217,43 +248,52 @@ class Mod():
         output = f'{os.path.join(Extract_Directory, self.hash)}'
         if not os.path.isdir(output):
             os.mkdir(output)
+        do_extract = False
+        on_disk = [os.path.join(root, file) for root, dirs, files in os.walk(output) for file in files]
+        for key in self.map:
+            file = os.path.join(output, self.raw_map[key])
+            if file not in on_disk:
+                do_extract = True
+                break
+        if do_extract:
             if self.type == ArchiveType.ZIP:
                 self.obj.extractall(path=output)
+            elif self.type == ArchiveType.RAR:
+                unrardll.extract(self.path, output)
             elif self.type == ArchiveType.SEVEN_ZIP:
                 self.obj.extractall(path=output)
         for key in self.map:
-            m.update(open(os.path.join(output, key), 'rb').read())
+            m.update(open(os.path.join(output, self.raw_map[key]), 'rb').read())
 
-    def write(self, old_list):
+    def write(self, old_list, list_file):
         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:
-                    copy = target[0] == '|';
-                    if copy:
-                        target = target[1:]
-                    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):
-                        if not copy:
-                            Error(f'Not replacing/removing non-symlink file: {target}')
-                        continue
-                    file = os.path.join(output, key)
-                    if copy:
-                        print(f'{file}\n (detached) -> {target}')
-                        shutil.copyfile(file, target)
-                    else:
-                        os.symlink(file, target)
+        for key in self.map:
+            targets = self.map[key]
+            for target in targets:
+                copy = target[0] == '|';
+                if copy:
+                    target = target[1:]
+                if target in old_list:
+                    old_list.remove(target)
+                list_file.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):
+                    if not copy:
+                        Warn(f'Not replacing/removing non-symlink file: {target}')
+                    continue
+                file = os.path.join(output, self.raw_map[key])
+                if copy:
+                    print(f'{file}\n (detached) -> {target}')
+                    shutil.copyfile(file, target)
+                else:
+                    os.symlink(file, target)
 
     def delete(self):
         for key in self.map:
@@ -263,7 +303,7 @@ class Mod():
                 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}')
+                    Warn(f'Not deleting non-symlink file: {target}')
 
 def is_armor_override(line):
     sp = line.split('/')
@@ -275,20 +315,62 @@ def is_armor_override(line):
         return True
     return False
 
+Seed = time.time()
+if Mode == 'write':
+    random.seed(Seed)
+
+with open(List_Path, 'a+') as l:
+    l.seek(0)
+    old_list = l.read().splitlines()
+    if len(old_list) > 2:
+        old_hash = old_list[0]
+        try:
+            old_seed = struct.unpack('!d', bytes.fromhex(old_list[1]))[0]
+        except ValueError:
+            Error(f'Invalid seed \'{old_list[1]}\' in {List_Path}.')
+        else:
+            if Mode == 'verify':
+                random.seed(old_seed)
+        old_list = old_list[2:]
+    if Mode == 'write':
+        l.truncate(0)
+
 Mods = []
+Shuffle = None
 Current_Mod = None
 Current_Override = None
 Global_Overrides = {}
 with open(Slots, 'r') as f:
-    for line in f.read().splitlines():
+    for ln, line in enumerate(f.read().splitlines()):
         if len(line) == 0 or line[0] == '#':
             continue
         if line == 'XXX':
             break
-        if line[0] == ';':
+        if line[0] in [';', 's']:
             if Current_Mod:
                 Mods.append(Current_Mod)
                 Current_Mod = None
+        if line[0] == 's':
+            Shuffle = []
+            path = os.path.join(Mods_Directory, line[1:])
+            if os.path.isfile(path):
+                Shuffle.append(path)
+            else:
+                Error(f'File not found: {path}')
+            continue
+        elif line[0] == '+':
+            path = os.path.join(Mods_Directory, line[1:])
+            if os.path.isfile(path):
+                Shuffle.append(path)
+            else:
+                Error(f'File not found: {path}')
+            continue
+        elif line[0:4].upper() == 'ROLL':
+            Current_Mod = Mod(random.choice(Shuffle))
+            Current_Override = None
+            Shuffle = None
+            continue
+        if line[0] == ';':
             path = os.path.join(Mods_Directory, line[1:])
             if os.path.isfile(path):
                 Current_Mod = Mod(path)
@@ -296,7 +378,7 @@ with open(Slots, 'r') as f:
             else:
                 Error(f'File not found: {path}')
             continue
-        cmd = ''
+        cmd = None
         if line[0] == '&':
             cmd = 'root'
         elif line[0] == '<':
@@ -313,16 +395,21 @@ with open(Slots, 'r') as f:
             line = line[1:]
             if line.startswith(NATIVEPC_SUBDIRS):
                 line = 'nativePC/' + line
+        else:
+            continue
         if cmd == 'root' and Current_Mod:
             Current_Mod.roots.append(line)
         elif cmd == 'ignore' and Current_Mod:
             Current_Mod.ignored.append(line)
         elif cmd == 'detach' and Current_Mod:
-            Current_Mod.ignored.append(line)
-            Current_Mod.map[line] = [f'|{line}']
+            if line not in Current_Mod.files:
+                Error(f'File not found: {Current_Mod.name}:{line}')
+            else:
+                Current_Mod.ignored.append(line)
+                Current_Mod.map[line] = [f'|{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}')
+                Error(f'File not found: {Current_Mod.name}:{line}')
             else:
                 Current_Override = line
                 Current_Mod.overrides[Current_Override] = []
@@ -353,29 +440,23 @@ for mod in Mods:
     mod.do_map(All_Files, Global_Overrides)
     mod.add_to_hash(M)
 
-with open(List_Path, 'a+') as l:
-    l.seek(0)
-    old_list = l.read().splitlines()
-    if Mode == 'write':
-        l.truncate(0)
-
 if Mode == 'write':
-    for mod in Mods:
-        mod.write(old_list)
+    with open(List_Path, 'a+') as l:
+        l.write(M.hexdigest() + '\n')
+        l.write(hex(struct.unpack('!Q', struct.pack('!d', Seed))[0])[2:] + '\n')
+        for mod in Mods:
+            mod.write(old_list, l)
     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.')
+    if old_hash == M.hexdigest():
+        print('Verification passed!')
+    else:
+        print('Verification failed.')
 elif Mode == 'delete':
     for mod in Mods:
         mod.delete()