A tool for setting up BI's P drive

setup_p_drive.py

#!/usr/bin/env python3

# Copyright 2020
# Author: razaqq <github.com/razaqq>

import os
from argparse import ArgumentParser
from os import walk, remove, mkdir, access
from os.path import join, exists, isfile, splitext
from shutil import rmtree, copy
from subprocess import Popen, PIPE
import sys
from winreg import OpenKey, QueryValueEx, CloseKey, KEY_READ, HKEY_LOCAL_MACHINE
from typing import List, Tuple


def print_err(text: str) -> None:
    print(f'\033[91m{text}\033[0m')


def print_warn(text: str) -> None:
    print(f'\033[93m{text}\033[0m')


def print_success(text: str) -> None:
    print(f'\033[92m{text}\033[0m')


def read_registry_key(prefix: int, path, key_name):
    root_key = OpenKey(prefix, path, 0, KEY_READ)
    value, _ = QueryValueEx(root_key, key_name)
    CloseKey(root_key)
    return value


def get_arma_paths() -> Tuple[str, str]:
    print('Getting ArmA paths...')

    base_paths = [
        r'SOFTWARE\Wow6432Node\Bohemia Interactive Studio',
        r'SOFTWARE\Bohemia Interactive Studio',
    ]

    arma2, arma2_found = '', False
    arma2oa, arma2oa_found = '', False

    for game in ('ArmA 2', 'ArmA 2 OA'):
        for base_path in base_paths:
            base_path += f'\\{game}'
            try:
                value = read_registry_key(HKEY_LOCAL_MACHINE, base_path, 'MAIN')
                if value and game == 'ArmA 2':
                    arma2 = value
                    break
                elif value and game == 'ArmA 2 OA':
                    arma2oa = value
                    break
            except FileNotFoundError:
                continue

    if not arma2:
        print_err('Cannot find arma 2 reg key')
        os.system('pause')
        sys.exit(1)

    print_success(f'Found ArmA 2 in: {arma2}')
    if not arma2oa:
        print_warn('Cannot find arma 2 oa path, proceeding...')
    else:
        print_success(f'Found ArmA 2 OA in: {arma2oa}')

    return arma2, arma2oa


def get_file_list_recursive(path: str, extension: str = '') -> List[str]:
    all_files = []
    for root, dirs, files in walk(path):
        for file in files:
            if file.endswith(extension):
                all_files.append(join(root, file))
    return all_files


def get_addon_list(arma2: str, arma2oa: str) -> List[str]:
    print('Getting addon list...')

    addons = []
    common = []
    expansion = []

    # TODO: what about betas, idk

    if arma2oa:
        common.extend(get_file_list_recursive(join(arma2oa, 'Common'), '.pbo'))
        expansion.extend(get_file_list_recursive(join(arma2oa, 'Expansion'), '.pbo'))

        if os.path.exists(join(arma2oa, 'Addons')):
            addons.extend(get_file_list_recursive(join(arma2oa, 'Addons'), '.pbo'))
        if os.path.exists(join(arma2oa, 'Dta')):
            addons.extend(get_file_list_recursive(join(arma2oa, 'Dta'), '.pbo'))
        if os.path.exists(join(arma2oa, '@RTEditor')):
            addons.extend(get_file_list_recursive(join(arma2oa, '@RTEditor'), '.pbo'))

    if arma2 != arma2oa:
        addons.extend(get_file_list_recursive(join(arma2, 'Addons'), '.pbo'))
        addons.extend(get_file_list_recursive(join(arma2, 'Dta'), '.pbo'))
        addons.extend(get_file_list_recursive(join(arma2, '@RTEditor'), '.pbo'))

    common.extend(addons)
    common.extend(expansion)

    print_success(f"Found {len(common)} pbo's to extract.")

    return common


def clear_p_drive() -> None:

    print('Clearing old P: drive...')

    try:
        if exists('p:\\ca'):
            rmtree('p:\\ca')
        if exists('p:\\languagecore'):
            rmtree('p:\\languagecore')
        if exists('p:\\core'):
            rmtree('p:\\core')
        if exists('p:\\bin'):
            rmtree('p:\\bin')
        if exists('p:\\ION_RTE'):
            rmtree('p:\\ION_RTE')
        if exists('p:\\ION_RTE_dta'):
            rmtree('p:\\ION_RTE_dta')
        if exists('p:\\dta'):
            rmtree('p:\\dta')
        if exists('p:\\WRP_PROJECTS'):
            rmtree('p:\\WRP_PROJECTS')
        if exists('p:\\corepatch_a2'):
            rmtree('p:\\corepatch_a2')
        if exists('p:\\buldozer.exe'):
            remove('p:\\buldozer.exe')
    except Exception as e:
        print_err(f'Failed to clear p: drive: {e}')
        sys.exit(1)

    print_success('P: drive clearing successful.')


def extract_pbos(addons: List[str]) -> None:
    print("Extracting PBO's (this will take a while)...")

    for addon in addons:
        p = Popen(['extractpbo', '-ap', addon, 'p:\\'], stderr=PIPE, stdout=PIPE)
        stdout, stderr = p.communicate()
        print(stderr.decode())

        if p.returncode != 0:
            print_err(f'Failed to extractpbo: {addon}')
            sys.exit(1)

    print_success('PBO extraction successful.')


def setup_buldozer(arma2oa: str) -> None:
    print('Setting up boldozer...')

    dta = join('p:', 'dta')
    if not exists(dta):
        mkdir(dta)

    copy(join(arma2oa, 'Expansion', 'dta', 'product.bin'), join(dta))
    copy(join(arma2oa, 'arma2oa.exe'), join('p:', 'buldozer.exe'))

    print_success('Buldozer setup successful.')


def paa_to_tga(convert_bin: str) -> None:
    print('Converting all paa to tga...')

    for paa in get_file_list_recursive('p:\\', '.paa'):
        p = Popen([convert_bin, paa, f'{splitext(paa)[0]}.tga'], stdout=PIPE, stderr=PIPE)
        stdout, stderr = p.communicate()
        print(stdout.decode().strip())
        if p.returncode != 0:
            print_err(f'Failed to convert paa: {paa}')

    print_success('Successfully converted all paa to tga.')


if __name__ == '__main__':
    parser = ArgumentParser(description='Setup tool for BI P: drive.')
    parser.add_argument('--clear', dest='clear', action='store_true', default=False,
                        help='Clears old stuff from p: - default: false')
    parser.add_argument('--pbo', dest='pbo', action='store_true', default=False,
                        help='Extracts pbos to p: - default: true')
    parser.add_argument('--buldozer', dest='buldozer', action='store_true', default=False,
                        help='Sets up boldozer in p: - default: true')
    parser.add_argument('--tga', dest='tga', action='store_true', default=False,
                        help='Converts all paa files in p: to tga - default: false')
    parser.add_argument('--pal2pace', dest='pal2pace', type=str, default=False)
    args = parser.parse_args()

    if True not in [v is True for _, v in args.__dict__.items()]:
        args.pbo = True
        args.buldozer = True

    if args.tga and not args.pal2pace:
        print_err('You need to specify the path to Pal2PacE.exe with --pal2pace when using --tga.')
        sys.exit(1)
    if args.tga and not (isfile(args.pal2pace) and access(args.pal2pace, os.X_OK)):
        print_err('The path to Pal2PacE.exe you specified does not seem to work. '
                      'Please download ArmA 3 Tools from steam or get BI Tools 2.5!')
        sys.exit(1)

    if not exists('p:\\'):
        print_err('No P: drive set')
        os.system('pause')
        sys.exit(1)

    arma2, arma2oa = get_arma_paths()

    if args.clear:
        clear_p_drive()
    if args.pbo:
        addons = get_addon_list(arma2, arma2oa)
        extract_pbos(addons)
    if args.buldozer:
        setup_buldozer(arma2oa)
    if args.tga:
        paa_to_tga(args.pal2pace)

    print_success('DONE!')