LibRay/libray/core.py

# -*- coding: utf8 -*-

# libray - Libre Blu-Ray PS3 ISO Tool
# Copyright © 2018 - 2024 Nichlas Severinsen
#
# This file is part of libray.
#
# libray 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.
#
# libray 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 libray.  If not, see <https://www.gnu.org/licenses/>.


import os
import sys
import stat
import zlib
import shutil
import threading
import requests
from bs4 import BeautifulSoup
from typing import Any, Optional


try:
    from libray import iso
    from libray import ird
except ImportError:
    import iso
    import ird

# Magic numbers / Constant variables

SECTOR = 2048
ALL_IRD_NET_LOC = 'http://jonnysp.bplaced.net/data.php'
GET_IRD_NET_LOC = 'http://jonnysp.bplaced.net/ird/'

# Utility functions


def to_int(data: bytes, byteorder: str = 'big') -> int:
    """Convert bytes to integer"""
    if isinstance(data, bytes):
        return int.from_bytes(data, byteorder)
    return 0


def to_bytes(data: str) -> Optional[bytes]:
    """Convert a string of HEX to bytes"""
    if isinstance(data, str):
        return bytes(bytearray.fromhex(data))
    return None


ISO_SECRET = to_bytes("380bcf0b53455b3c7817ab4fa3ba90ed")
ISO_IV = to_bytes("69474772af6fdab342743aefaa186287")


def size(path: str) -> int:
    """Get size of a file or block device in bytes"""
    pathstat = os.stat(path)

    # Check if it's a block device

    if stat.S_ISBLK(pathstat.st_mode):
        return open(path, 'rb').seek(0, os.SEEK_END)

    # Otherwise, it's hopefully a file

    return pathstat.st_size


def read_seven_bit_encoded_int(fileobj, order: str) -> int:
    """Read an Int32, 7 bits at a time."""
    # The highest bit of the byte, when on, means to continue reading more bytes.
    count = 0
    shift = 0
    byte = -1
    while (byte & 0x80) != 0 or byte == -1:
        # Check for a corrupted stream. Read a max of 5 bytes.
        if shift == (5 * 7):
            raise ValueError
        byte = to_int(fileobj.read(1), order)
        count |= (byte & 0x7F) << shift
        shift += 7
    return count


def error(msg: str) -> None:
    """Print fatal error message and terminate"""
    print(f'[ERROR] {msg}', file=sys.stderr)
    sys.exit(1)


def warning(msg: str, args: Any) -> None:
    """Print a warning message. Warning messages can be silenced with --quiet"""

    if not args.quiet:
        print(f'[WARNING] {msg}. Continuing regardless.', file=sys.stderr)


def vprint(msg: str, args: Any) -> None:
    """Vprint, verbose print, can be silenced with --quiet"""

    if not args.quiet:
        print(f'[*] {msg}')


def download_ird(ird_name: str) -> None:
    """Download an .ird from GET_IRD_NET_LOC"""

    # Check if file already exists and skip if it does
    if os.path.exists(ird_name):  # TODO: might want to check that the file is valid first, could do a HEAD agains the url
        return

    ird_link = GET_IRD_NET_LOC + ird_name
    r = requests.get(ird_link, stream=True, timeout=30)

    with open(ird_name, 'wb') as ird_file:
        r.raw.decode_content = True
        shutil.copyfileobj(r.raw, ird_file)


def ird_by_game_id(game_id: str) -> Optional[str]:
    """Using a game_id, download the responding .ird from ALL_IRD_NET_LOC"""
    gameid = game_id.replace('-', '')
    try:
        r = requests.get(ALL_IRD_NET_LOC, headers={'User-Agent': 'Anonymous (You)'}, timeout=5)
    except requests.exceptions.ReadTimeout:
        error('Server timed out, fix your connection or manually specify a key/ird.')
    soup = BeautifulSoup(r.text, "html.parser")

    ird_name = False
    for elem in soup.find_all("a"):
        url = elem.get('href').split('/')[-1].replace('\\"', '')
        if gameid in url:
            ird_name = url

    if not ird_name:
        error("Unable to download IRD, couldn't find link. You could specify the decryption key with -d if you have it.")

    download_ird(ird_name)

    return (ird_name)


def crc32(filename: str, cancel: Optional[threading.Event] = None) -> Optional[str]:
    """Calculate crc32 for file.

    If *cancel* is provided, the computation can be aborted from another
    thread by calling ``cancel.set()``.  Returns ``None`` when cancelled.
    """
    if cancel is None:
        cancel = threading.Event()

    with open(filename, 'rb') as infile:
        crc_val = 0
        while not cancel.is_set():
            data = infile.read(65536)
            if not data:
                break
            crc_val = zlib.crc32(data, crc_val)

        if cancel.is_set():
            return None

        return f"{crc_val & 0xFFFFFFFF:08X}"


def serial_country(title: str) -> str:
    """Get country from disc serial / productcode / title_id"""

    if title[2] == 'A':
        return 'Asia'
    if title[2] == 'C':
        return 'China'
    if title[2] == 'E':
        return 'Europe'
    if title[2] == 'H':
        return 'Hong Kong'
    if title[2] == 'J' or title[2] == 'P':
        return 'Japan'
    if title[2] == 'K':
        return 'Korea'
    if title[2] == 'U' or title[2] == 'T':
        return 'USA'

    raise ValueError('Unknown country?!')


def multiman_title(title: str) -> str:
    """Fix special characters in title for Multiman style"""

    replace = {
        ':': ' -',
        '/': '-',
        '™': '',
        '®': '',
    }

    for key, val in replace.items():
        title = title.replace(key, val)

    return title


# Main functions

def info(args: Any) -> None:
    """Print information about .iso and then quit."""

    if args.iso:
        input_iso = iso.ISO(args)
        input_iso.print_info()
        sys.exit()

    if args.ird:
        input_ird = ird.IRD(args.ird)
        input_ird.print_info(regions=True)
        sys.exit()


def decrypt(args: Any) -> None:
    """Try to decrypt a given .iso using relevant .ird or encryption key from argparse

    If no .ird is given this will try to automatically download an .ird file with the encryption/decryption key for the given game .iso
    """

    input_iso = iso.ISO(args)

    # TODO: some of the logic should probably be moved up here instead of residing in the decrypt function
    input_iso.decrypt(args)


def encrypt(args: Any) -> None:
    """Try to re-encrypt a decrypted .iso using relevant .ird or encryption key from argparse

    If no .ird is given this will try to automatically download an .ird file with the encryption/decryption key for the given game .iso
    """

    input_iso = iso.ISO(args)

    input_iso.encrypt(args)