#!/bin/bash
#  ____  _   _                      
# / ___|| |_| |__   ___  _ __   ___ 
# \___ \| __| '_ \ / _ \| '_ \ / _ \
#  ___) | |_| | | | (_) | |_) |  __/
# |____/ \__|_| |_|\___/| .__/ \___|
#                       |_|         
#
#

TODO_FOLDER="./todo"
DKEYS_FOLDER="./dkeys"

DONE_FOLDER="./done"
ENCRYPTED_ISO="$DONE_FOLDER/encrypted"
DECRYPTED_ISO="$DONE_FOLDER/decrypted"
EXTRACTED_ISO="$DONE_FOLDER/extracted"
ZIP_FOLDER="$DONE_FOLDER/zip"
NEED_KEY="$DONE_FOLDER/need_key"
TO_DELETE_FOLDER="$DONE_FOLDER/to_delete"

install_dependencies() {
    if [ -x "$(command -v apt-get)" ]; then
        echo "Detected Debian-based distribution."
        sudo apt-get update
        sudo apt-get install -y p7zip-full wget
        sudo pip install libray
    elif [ -x "$(command -v yum)" ]; then
        echo "Detected Red Hat-based distribution."
        sudo yum install -y p7zip wget
        sudo pip install libray
    elif [ -x "$(command -v pacman)" ]; then
        echo "Detected Arch-based distribution."
        sudo pacman -Sy --noconfirm p7zip wget
        sudo pip install libray
    elif [ -x "$(command -v dnf)" ]; then
        echo "Detected Rocky Linux."
        sudo dnf install -y p7zip wget
        sudo pip install libray
    else
        echo "Unsupported distribution. Please install Python 3 manually."
        exit 1
    fi
}

check_if_folders_exists() {
    folders=("$ENCRYPTED_ISO" "$DECRYPTED_ISO" "$EXTRACTED_ISO" "$ZIP_FOLDER" "$NEED_KEY" "$TO_DELETE_FOLDER")
    for folder in "${folders[@]}"; do
        if [ ! -d "$folder" ]; then
            mkdir -p "$folder"
        fi
    done
}

get_dkeys() {
    wget -q https://git.sthope.dev/sthope/Emulation-and-ROMs/raw/branch/main/PS3/PS3roms/dkeys.7z -P ./
    7z x ./dkeys.7z >/dev/null 2>&1
    rm -rf ./dkeys.7z
}

extract_zips() {
    for file in "$TODO_FOLDER"/*.zip; do
        if [ -f "$file" ]; then
            7z x "$file" -o"$ENCRYPTED_ISO">/dev/null 2>&1
            if [ $? -ne 0 ]; then
                echo "Error extracting $file"
            fi
        fi
    done
}

decrypt_iso() {
    for file in "$ENCRYPTED_ISO"/*.iso; do
        if [ -f "$file" ]; then
            decrypted_iso="$DECRYPTED_ISO/$(basename "$file")"
            if libray -i "$file" -o "$decrypted_iso" >/dev/null 2>&1; then
                mv "$file" "$TO_DELETE_FOLDER"
            else
                mv "$file" "$NEED_KEY"
                for key_file in "$NEED_KEY"/*.iso; do
                    if [ -f "$key_file" ]; then
                        get_dkeys
                        decrypted_iso="$DECRYPTED_ISO/$(basename "$key_file")"
                        dkey_file="$DKEYS_FOLDER/$(basename "$key_file" .iso).dkey"
                        DKEY="$(cat "$dkey_file")"
                        libray -i "$key_file" -o "$decrypted_iso" -d "$DKEY" >/dev/null 2>&1
                        if [ $? -eq 0 ]; then
                            mv "$key_file" "$TO_DELETE_FOLDER"
                        else
                            echo "Error decrypting $key_file"
                        fi
                    fi
                done               
            fi
        fi
    done
    rm -rf $DKEYS_FOLDER
}

extract_decrypted() {
    for file in "$DECRYPTED_ISO"/*.iso; do
        if [ -f "$file" ]; then
            decrypted_iso=$(basename "$DECRYPTED_ISO/$(basename "$file" ".$extension" | cut -d. -f1)")
            7z x "$file" -o"$EXTRACTED_ISO/$decrypted_iso">/dev/null 2>&1
            if [ $? -ne 0 ]; then
                echo "Error extracting decrypted $file"
            fi
        fi
    done
}

decrypt_with_dkeys() {
    for file in "$NEED_KEY"/*.iso; do
        if [ -f "$file" ]; then
            decrypted_iso="$DECRYPTED_ISO/$(basename "$file")"
            key_file="$DKEYS_FOLDER/$(basename "$file" .iso).dkey"
            DKEY="$(cat "$key_file")"
            libray -i "$file" -o "$decrypted_iso" -d "$DKEY"
            if [ $? -eq 0 ]; then
                mv "$file" "$TO_DELETE_FOLDER"
            else
                echo "Error extracting $file"
            fi
        fi
    done
}

cleanup() {
    case "$1" in
        --decrypted_isos)
            mv "$DECRYPTED_ISO" ./
            rm -rf "$DONE_FOLDER"
            ;;
        --extracted_isos)
            mv "$EXTRACTED_ISO" ./
            rm -rf "$DONE_FOLDER"
            ;;
        --zip_folder)
            mv "$ZIP_FOLDER" ./
            rm -rf "$DONE_FOLDER"
            ;;
        --encrypted_isos)
            mv "$TO_DELETE_FOLDER" ./
            rm -rf "$DONE_FOLDER"
            ;;
        *)
            echo "Invalid option. Usage: $0 --clean-up [--decrypted_isos|--extracted_isos|--zip_folder|--encrypted_isos]"
            exit 1
            ;;
    esac
}

if [ "$1" == "--extract-all" ]; then
    check_if_folders_exists
    extract_zips
elif [ "$1" == "--decrypt-all" ]; then
    check_if_folders_exists
    decrypt_iso
elif [ "$1" == "--extract-decrypted-all" ]; then
    check_if_folders_exists
    extract_decrypted
elif [ "$1" == "--need-dkey" ]; then
    check_if_folders_exists
    check_keys
elif [ "$1" == "--clean-up" ]; then
    cleanup $2
elif [ "$1" == "--all" ]; then
    check_if_folders_exists
    extract_zips
    decrypt_iso
    extract_decrypted
    cleanup $2
else
    echo "Usage: $0 {--extract-all|--decrypt-all|--extract-decrypted-all|--need-dkey|--all}"
    exit 1
fi