TinGen.py

#!/usr/bin/env python3
# -*- coding: utf8 -*-
from google.auth.credentials import Credentials
from TinGen.utils import create_tinfoil_index
from TinGen.utils import CompressionFlag
from argparse import ArgumentParser
from TinGen import TinGen
from pathlib import Path

if __name__ == '__main__':
    parser = ArgumentParser(
        description='Script that will allow you to generate an index file ' +
        'with Google Drive file links for use with Tinfoil',
    )
    parser.add_argument(
        nargs='*',
        metavar='FOLDER_ID_TO_SCAN',
        dest='folder_ids',
        help='Folder IDs of Google Drive folders to scan',
    )
    parser.add_argument(
        '--credentials',
        default='credentials.json',
        metavar='CREDENTIALS_FILE_NAME',
        help='Path to Google Application Credentials',
    )
    parser.add_argument(
        '--token',
        default='gdrive.token',
        metavar='TOKEN_FILE_PATH',
        help='Path to Google OAuth2.0 User Token',
    )
    parser.add_argument(
        '--headless',
        action='store_true',
        help='Allows to perform Google OAuth2.0 User Token Authentication ' +
        'in headless environment',
    )
    parser.add_argument(
        '--index-file',
        metavar='INDEX_FILE_PATH',
        default='index.tfl',
        help='Path to output index file',
    )
    parser.add_argument(
        '--share-files',
        action='store_true',
        help='Share all files inside the index file',
    )
    parser.add_argument(
        '--no-recursion',
        dest='recursion',
        action='store_false',
        help='Scans for files only in top directory for each Folder ID ' +
        'entered',
    )
    parser.add_argument(
        '--add-nsw-files-without-title-id',
        action='store_true',
        help='Adds files without valid Title ID',
    )
    parser.add_argument(
        '--add-non-nsw-files',
        action='store_true',
        help='Adds files without valid NSW ROM extension(NSP/NSZ/XCI/XCZ) ' +
        'to index',
    )
    parser.add_argument(
        '--add-nsw-info-to-success',
        action='store_true',
        help='Adds simple information of NSW titles added to index',
    )
    parser.add_argument(
        '--add-update-date-to-success',
        action='store_true',
        help='Adds index update date to index',
    )
    parser.add_argument(
        '--add-update-time-to-success',
        action='store_true',
        help='Adds index update time to index',
    )
    parser.add_argument(
        '--success',
        metavar='SUCCESS_MESSAGE',
        help='Adds a success message to index file to show if index is ' +
        'successfully read by Tinfoil',
    )
    parser.add_argument(
        '--encrypt',
        action='store_true',
        help='Encrypts the resulting index file with AES-ECB-256',
    )
    parser.add_argument(
        '--public-key',
        metavar='PUBLIC_KEY_FILE_PATH',
        default='public.pem',
        help='Path to RSA Public Key to encrypt AES-ECB-256 key with',
    )
    parser.add_argument(
        '--vm-file',
        help='Path to VM File',
    )
    parser.add_argument(
        '--upload-to-folder-id',
        metavar='UPLOAD_FOLDER_ID',
        dest='upload_folder_id',
        help='Upload resulting index to Folder ID supplied',
    )
    parser.add_argument(
        '--upload-to-my-drive',
        action='store_true',
        help='Upload resulting index to My Drive',
    )
    parser.add_argument(
        '--new-upload-id',
        action='store_true',
        help='Uploads the newly generated index file with a new File ID ' +
        'instead of replacing old one',
    )
    parser.add_argument(
        '--share-uploaded-index',
        action='store_true',
        help='Shares the index file that is uploaded to Google Drive',
    )
    parser.add_argument(
        '--tinfoil-min-ver',
        metavar='TINFOIL_MINIMUM_VERSION',
        default='7.00',
        type=str,
        help='Minimum Tinfoil client version to use index with',
    )

    task_parser = parser.add_mutually_exclusive_group()
    task_parser.add_argument(
        '--auth',
        action='store_true',
        help='Run Google User Token authorize task if token doesn\'t exist',
    )
    task_parser.add_argument(
        '--generator',
        action='store_true',
        help='Run index generation task',
    )

    compression_parser = parser.add_mutually_exclusive_group()
    compression_parser.add_argument(
        '--zstandard',
        '--zstd',
        action='store_true',
        help='Compresses index with Zstandard compression method',
    )
    compression_parser.add_argument(
        '--zlib',
        action='store_true',
        help='Compresses index with zlib compression method',
    )
    compression_parser.add_argument(
        '--no-compress',
        action='store_true',
        help='Flag to not compress index',
    )

    theme_opts_parser = parser.add_argument_group()
    theme_opts_parser.add_argument(
        '--theme-blacklist',
        nargs='*',
        help='Theme IDs to add to index to blacklist',
    )
    theme_opts_parser.add_argument(
        '--theme-whitelist',
        nargs='*',
        help='Theme IDs to add to index to whitelist',
    )
    theme_opts_parser.add_argument(
        '--theme-error',
        metavar='ERROR_MESSAGE',
        help='Error message to show if theme check fails',
    )

    args = parser.parse_args()

    theme_msg = None
    if args.theme_error:
        theme_msg = args.theme_error.replace('\\n', '\n').replace('\\t', '\t')

    theme_blacklist = None
    if args.theme_blacklist:
        theme_blacklist = args.theme_blacklist

    theme_whitelist = None
    if args.theme_whitelist:
        theme_whitelist = args.theme_whitelist

    generator = TinGen(
        args.token,
        args.credentials,
        args.headless,
        args.tinfoil_min_ver,
        theme_blacklist=theme_blacklist,
        theme_whitelist=theme_whitelist,
        theme_error=theme_msg,
    )

    if args.auth:
        credentials = generator.gdrive_service._get_creds(
            credentials=args.credentials,
            token=args.token,
            headless=args.headless,
        )
        if isinstance(credentials, Credentials):
            print('Token generated successfully!')
        else:
            raise RuntimeError('Unable to generate OAuth2 user credentials. ' +
                               'Unable to continue!')

    else:
        print('Generating index')
        generator.index_generator(
            args.folder_ids,
            args.recursion,
            args.add_nsw_files_without_title_id,
            args.add_non_nsw_files,
        )

        if args.add_nsw_info_to_success:
            print('Adding NSW title information message to index')
            generator.add_nsw_title_info_to_success()

        if args.add_update_date_to_success or args.add_update_time_to_success:
            print('Adding date/time information to index')
            generator.add_datetime_to_success(
                args.add_update_date_to_success,
                args.add_update_time_to_success,
            )

        if args.success:
            print('Adding success message to index')
            generator.update_index_success_message(
                args.success.replace('\\n', '\n').replace('\\t', '\t'),
            )

        compression_flag = CompressionFlag.ZSTD_COMPRESSION

        if args.zstandard:
            compression_flag = CompressionFlag.ZSTD_COMPRESSION
        elif args.zlib:
            compression_flag = CompressionFlag.ZLIB_COMPRESSION
        elif args.no_compress:
            compression_flag = CompressionFlag.NO_COMPRESSION

        print(f'Creating generated index to {args.index_file}')
        if args.encrypt:
            rsa_pub_key_path = None
            if args.public_key:
                rsa_pub_key_path = Path(args.public_key)

            vm_path = None
            if args.vm_file:
                vm_path = Path(args.vm_file)

            create_tinfoil_index(
                generator.index,
                Path(args.index_file),
                compression_flag,
                rsa_pub_key_path=rsa_pub_key_path,
                vm_path=vm_path,
            )
        else:
            create_tinfoil_index(
                generator.index,
                Path(args.index_file),
                compression_flag,
            )

        if args.share_files:
            print('Sharing files in index')
            for folder_id in args.folder_ids:
                generator.gdrive_service.share_file(folder_id)
            # generator.share_index_files()

        if args.upload_folder_id:
            print(f'Uploading {args.index_file} to {args.upload_folder_id}')
            generator.gdrive_service.upload_file(
                args.index_file,
                args.upload_folder_id,
                args.share_uploaded_index,
                args.new_upload_id,
            )

        if args.upload_to_my_drive:
            print(f'Uploading {args.index_file} to \"My Drive\"')
            generator.gdrive_service.upload_file(
                args.index_file,
                None,
                args.share_uploaded_index,
                args.new_upload_id,
            )

        print('Index Generation Complete')