Initial parallelization attempt
This commit is contained in:
parent
12c05afa57
commit
35e49044f1
SSH key fingerprint:
SHA256:x4/RtnjUyuHkdvmwNDsWSfcfF1V5PNr3OpriZqOvCX8
10 changed files with 321 additions and 143 deletions
278
libray/iso.py
278
libray/iso.py
|
|
@ -23,9 +23,12 @@ import os
|
|||
import sys
|
||||
import sqlite3
|
||||
import pathlib
|
||||
from threading import Thread
|
||||
import concurrent.futures
|
||||
import threading
|
||||
import time
|
||||
import pkg_resources
|
||||
import queue
|
||||
from threading import Thread
|
||||
from importlib import resources
|
||||
from tqdm import tqdm
|
||||
from Crypto.Cipher import AES
|
||||
|
||||
|
|
@ -40,6 +43,28 @@ except ImportError:
|
|||
import sfo
|
||||
|
||||
|
||||
def _decrypt_sector_worker(disc_key, sector_data, sector_number):
|
||||
"""Standalone worker for parallel sector decryption."""
|
||||
iv = bytearray(16)
|
||||
num = sector_number
|
||||
for j in range(16):
|
||||
iv[15 - j] = num & 0xFF
|
||||
num >>= 8
|
||||
cipher = AES.new(disc_key, AES.MODE_CBC, bytes(iv))
|
||||
return (sector_number, cipher.decrypt(sector_data))
|
||||
|
||||
|
||||
def _encrypt_sector_worker(disc_key, sector_data, sector_number):
|
||||
"""Standalone worker for parallel sector encryption."""
|
||||
iv = bytearray(16)
|
||||
num = sector_number
|
||||
for j in range(16):
|
||||
iv[15 - j] = num & 0xFF
|
||||
num >>= 8
|
||||
cipher = AES.new(disc_key, AES.MODE_CBC, bytes(iv))
|
||||
return (sector_number, cipher.encrypt(sector_data))
|
||||
|
||||
|
||||
class ISO:
|
||||
"""Class for handling PS3 .iso files.
|
||||
|
||||
|
|
@ -166,127 +191,182 @@ class ISO:
|
|||
if args.verbose and not args.quiet:
|
||||
self.print_info()
|
||||
|
||||
def _make_iv(self, sector_number):
|
||||
"""Build a 16-byte IV from a sector number (little-endian)."""
|
||||
iv = bytearray(16)
|
||||
num = sector_number
|
||||
for j in range(16):
|
||||
iv[15 - j] = num & 0xFF
|
||||
num >>= 8
|
||||
return bytes(iv)
|
||||
|
||||
def _process_region_pipeline(self, input_path, region, num_workers, encrypt_mode, args):
|
||||
"""Process an encrypted region using a reader-worker pipeline.
|
||||
|
||||
A reader thread reads sectors from the file and puts them on a queue.
|
||||
Worker threads pull from the queue, process each sector in parallel,
|
||||
and store results. This overlaps I/O with processing for better CPU usage.
|
||||
"""
|
||||
num_sectors = (region['end'] - region['start']) // core.SECTOR
|
||||
queue_size = min(64, num_sectors)
|
||||
sector_queue = queue.Queue(maxsize=queue_size)
|
||||
results = [None] * num_sectors
|
||||
results_lock = threading.Lock()
|
||||
|
||||
def reader():
|
||||
with open(input_path, 'rb') as f:
|
||||
f.seek(region['start'])
|
||||
for i in range(num_sectors):
|
||||
sector_data = f.read(core.SECTOR)
|
||||
sector_queue.put((i, sector_data))
|
||||
for _ in range(num_workers):
|
||||
sector_queue.put(None)
|
||||
|
||||
def worker():
|
||||
while True:
|
||||
item = sector_queue.get()
|
||||
if item is None:
|
||||
sector_queue.task_done()
|
||||
break
|
||||
idx, sector_data = item
|
||||
_, processed = (_encrypt_sector_worker if encrypt_mode else _decrypt_sector_worker)(
|
||||
self.disc_key, sector_data, region['start'] // core.SECTOR + idx
|
||||
)
|
||||
with results_lock:
|
||||
results[idx] = processed
|
||||
sector_queue.task_done()
|
||||
|
||||
reader_thread = threading.Thread(target=reader, daemon=True)
|
||||
reader_thread.start()
|
||||
|
||||
workers = []
|
||||
for _ in range(num_workers):
|
||||
t = threading.Thread(target=worker, daemon=True)
|
||||
t.start()
|
||||
workers.append(t)
|
||||
|
||||
for t in workers:
|
||||
t.join()
|
||||
|
||||
reader_thread.join()
|
||||
|
||||
return b''.join(results)
|
||||
|
||||
def decrypt(self, args):
|
||||
"""Decrypt self using args from argparse."""
|
||||
|
||||
core.vprint(f'Decrypting with disc key: {self.disc_key.hex()}', args)
|
||||
|
||||
with open(args.iso, 'rb') as input_iso:
|
||||
num_workers = args.threads if args.threads and args.threads > 0 else os.cpu_count() or 1
|
||||
if num_workers > 1:
|
||||
core.vprint(f'Using {num_workers} threads for parallel decryption', args)
|
||||
|
||||
if not args.output:
|
||||
output_name = f'{self.game_id}.iso'
|
||||
else:
|
||||
output_name = args.output
|
||||
if not args.output:
|
||||
output_name = f'{self.game_id}.iso'
|
||||
else:
|
||||
output_name = args.output
|
||||
|
||||
core.vprint(f'Decrypted .iso is output to: {output_name}', args)
|
||||
core.vprint(f'Decrypted .iso is output to: {output_name}', args)
|
||||
|
||||
with open(output_name, 'wb') as output_iso:
|
||||
total_sectors = self.size // core.SECTOR
|
||||
|
||||
if not args.quiet:
|
||||
pbar = tqdm(total=(self.size // 2048))
|
||||
with open(args.iso, 'rb') as input_iso, open(output_name, 'wb') as output_iso:
|
||||
|
||||
for region in self.regions:
|
||||
pbar = tqdm(total=total_sectors, file=sys.stdout, disable=args.quiet, leave=True)
|
||||
|
||||
for region in self.regions:
|
||||
region_sectors = (region['end'] - region['start']) // core.SECTOR
|
||||
|
||||
if not region['enc']:
|
||||
# Unencrypted region — copy sequentially
|
||||
input_iso.seek(region['start'])
|
||||
|
||||
# Unencrypted region, just copy it
|
||||
if not region['enc']:
|
||||
while input_iso.tell() < region['end']:
|
||||
data = input_iso.read(core.SECTOR)
|
||||
if not data:
|
||||
core.warning('Trying to read past the end of the file', args)
|
||||
break
|
||||
output_iso.write(data)
|
||||
|
||||
if not args.quiet:
|
||||
pbar.update(1)
|
||||
continue
|
||||
# Encrypted region, decrypt then write
|
||||
for _ in range(region_sectors):
|
||||
data = input_iso.read(core.SECTOR)
|
||||
if not data:
|
||||
core.warning('Trying to read past the end of the file', args)
|
||||
break
|
||||
output_iso.write(data)
|
||||
pbar.update(1)
|
||||
else:
|
||||
# Encrypted region — pipeline: reader thread + worker threads
|
||||
if num_workers > 1:
|
||||
processed = self._process_region_pipeline(
|
||||
args.iso, region, num_workers, encrypt_mode=False, args=args
|
||||
)
|
||||
else:
|
||||
while input_iso.tell() < region['end']:
|
||||
num = input_iso.tell() // 2048
|
||||
iv = bytearray([0 for i in range(0, 16)])
|
||||
for j in range(0, 16):
|
||||
iv[16 - j - 1] = (num & 0xFF)
|
||||
num >>= 8
|
||||
# Sequential fallback
|
||||
input_iso.seek(region['start'])
|
||||
processed = bytearray()
|
||||
for i in range(region_sectors):
|
||||
sector_num = region['start'] // core.SECTOR + i
|
||||
iv = self._make_iv(sector_num)
|
||||
cipher = AES.new(self.disc_key, AES.MODE_CBC, iv)
|
||||
processed.extend(cipher.decrypt(input_iso.read(core.SECTOR)))
|
||||
processed = bytes(processed)
|
||||
|
||||
data = input_iso.read(core.SECTOR)
|
||||
if not data:
|
||||
core.warning('Trying to read past the end of the file', args)
|
||||
break
|
||||
output_iso.write(processed)
|
||||
pbar.update(region_sectors)
|
||||
|
||||
cipher = AES.new(self.disc_key, AES.MODE_CBC, bytes(iv))
|
||||
decrypted = cipher.decrypt(data)
|
||||
|
||||
output_iso.write(decrypted)
|
||||
|
||||
if not args.quiet:
|
||||
pbar.update(1)
|
||||
|
||||
if not args.quiet:
|
||||
pbar.close()
|
||||
|
||||
core.vprint('Decryption complete!', args)
|
||||
pbar.close()
|
||||
core.vprint('Decryption complete!', args)
|
||||
|
||||
def encrypt(self, args):
|
||||
"""Encrypt self using args from argparse."""
|
||||
|
||||
core.vprint(f'Re-encrypting with disc key: {self.disc_key.hex()}', args)
|
||||
|
||||
with open(args.iso, 'rb') as input_iso:
|
||||
num_workers = args.threads if args.threads and args.threads > 0 else os.cpu_count() or 1
|
||||
if num_workers > 1:
|
||||
core.vprint(f'Using {num_workers} threads for parallel re-encryption', args)
|
||||
|
||||
if not args.output:
|
||||
output_name = f'{self.game_id}_e.iso'
|
||||
else:
|
||||
output_name = args.output
|
||||
if not args.output:
|
||||
output_name = f'{self.game_id}_e.iso'
|
||||
else:
|
||||
output_name = args.output
|
||||
|
||||
core.vprint(f'Re-encrypted .iso is output to: {output_name}', args)
|
||||
core.vprint(f'Re-encrypted .iso is output to: {output_name}', args)
|
||||
|
||||
with open(output_name, 'wb') as output_iso:
|
||||
with open(args.iso, 'rb') as input_iso, open(output_name, 'wb') as output_iso:
|
||||
|
||||
if not args.quiet:
|
||||
pbar = tqdm(total=(self.size // 2048))
|
||||
pbar = tqdm(total=(self.size // 2048), file=sys.stdout, disable=args.quiet, leave=True)
|
||||
|
||||
for region in self.regions:
|
||||
for region in self.regions:
|
||||
region_sectors = (region['end'] - region['start']) // core.SECTOR
|
||||
|
||||
if not region['enc']:
|
||||
# Unencrypted region — copy sequentially
|
||||
input_iso.seek(region['start'])
|
||||
|
||||
# Unencrypted region, just copy it
|
||||
if not region['enc']:
|
||||
while input_iso.tell() < region['end']:
|
||||
data = input_iso.read(core.SECTOR)
|
||||
if not data:
|
||||
core.warning('Trying to read past the end of the file', args)
|
||||
break
|
||||
output_iso.write(data)
|
||||
|
||||
if not args.quiet:
|
||||
pbar.update(1)
|
||||
continue
|
||||
# Decrypted region, re-encrypt it
|
||||
for _ in range(region_sectors):
|
||||
data = input_iso.read(core.SECTOR)
|
||||
if not data:
|
||||
core.warning('Trying to read past the end of the file', args)
|
||||
break
|
||||
output_iso.write(data)
|
||||
pbar.update(1)
|
||||
else:
|
||||
# Encrypted region — pipeline: reader thread + worker threads
|
||||
if num_workers > 1:
|
||||
processed = self._process_region_pipeline(
|
||||
args.iso, region, num_workers, encrypt_mode=True, args=args
|
||||
)
|
||||
else:
|
||||
while input_iso.tell() < region['end']:
|
||||
num = input_iso.tell() // 2048
|
||||
iv = bytearray([0 for i in range(0, 16)])
|
||||
for j in range(0, 16):
|
||||
iv[16 - j - 1] = (num & 0xFF)
|
||||
num >>= 8
|
||||
# Sequential fallback
|
||||
input_iso.seek(region['start'])
|
||||
processed = bytearray()
|
||||
for i in range(region_sectors):
|
||||
sector_num = region['start'] // core.SECTOR + i
|
||||
iv = self._make_iv(sector_num)
|
||||
cipher = AES.new(self.disc_key, AES.MODE_CBC, iv)
|
||||
processed.extend(cipher.encrypt(input_iso.read(core.SECTOR)))
|
||||
processed = bytes(processed)
|
||||
|
||||
data = input_iso.read(core.SECTOR)
|
||||
if not data:
|
||||
core.warning('Trying to read past the end of the file', args)
|
||||
break
|
||||
output_iso.write(processed)
|
||||
pbar.update(region_sectors)
|
||||
|
||||
cipher = AES.new(self.disc_key, AES.MODE_CBC, bytes(iv))
|
||||
encrypted = cipher.encrypt(data)
|
||||
if not args.quiet:
|
||||
pbar.close()
|
||||
|
||||
output_iso.write(encrypted)
|
||||
|
||||
if not args.quiet:
|
||||
pbar.update(1)
|
||||
|
||||
if not args.quiet:
|
||||
pbar.close()
|
||||
|
||||
core.vprint('Re-encryption complete!', args)
|
||||
core.vprint('Re-encryption complete!', args)
|
||||
|
||||
def get_key_from_args(self, game_title, args):
|
||||
# key provided with -d / --decryption-key
|
||||
|
|
@ -313,8 +393,16 @@ class ISO:
|
|||
core.vprint('Checking for bundled redump keys', args)
|
||||
|
||||
try:
|
||||
db = sqlite3.connect(pkg_resources.resource_filename(__name__, 'data/keys.db'))
|
||||
except FileNotFoundError:
|
||||
db_path = resources.files(__name__).joinpath('data', 'keys.db')
|
||||
if hasattr(db_path, 'read_bytes'):
|
||||
# importlib.resources.abc.Traversable - write to temp file for sqlite3
|
||||
import tempfile
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix='.db') as tmp:
|
||||
tmp.write(db_path.read_bytes())
|
||||
db = sqlite3.connect(tmp.name)
|
||||
else:
|
||||
db = sqlite3.connect(str(db_path))
|
||||
except (FileNotFoundError, AttributeError):
|
||||
db = sqlite3.connect((pathlib.Path(__file__).resolve() / 'data/') / 'keys.db')
|
||||
c = db.cursor()
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue