hide hard disk from WiiU (and use with vWii)

py_wiiu_hide_disk.py

#!/usr/bin/env python

"""
py_wiiu_hide_disk.py: hide hard disk from WiiU (and use with vWii)
Copyright (C) 2019 Vetter Michael

This program 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.

This program 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 this program.  If not, see <https://www.gnu.org/licenses/>.
"""

#
# this tools allow to change boot signature of hard disk to hide from WiiU (and use with vWii)
#
# thanks to https://code.google.com/archive/p/u-stealth/
#
# mbr = first 512 bytes
#
# last two bytes on bootable drives are usually 55 AA
# this tool changes either to 55AB (hidden) or to 55AA (shown)

import argparse
import os
import sys
import subprocess

# constants
STARTING_OFFSET=510 # second last byte of mbr (2 bytes are overwritten at offset 510 and 511)
HIDE_VAL=[0x55,0xab] # ["55","ab"] is also possible
SHOW_VAL=[0x55,0xaa]

def create_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--device",
                        help="device path (e.g. /dev/sdc)",
                        required=True
        )
    parser.add_argument("--hide",
                        help="hide drive",
                        action="store_const",
                        const=True
        )
    parser.add_argument("--unhide",
                        help="unhide drive",
                        action="store_const",
                        const=True
        )
    parser.add_argument("--skip-hwinfo",
                        help="skip hwinfo check",
                        action="store_const",
                        const=True
        )
    return(parser.parse_args())

def check_root():
    if os.getuid() != 0:
        print("please restart as root")
        exit(1)
    return

def check_answer(answer):
    if answer != "y" and answer != "Y":
        print("exiting")
        exit()

def check_dev(dev):
    # check for block device
    if os.st.S_ISBLK(os.stat(dev).st_mode) is False:
        answer = input("%s is not a block device. Are you sure this is the right path (y/Y)? " %(dev))
        check_answer(answer)
    # simple check if device path ends with number, which usually is a partiton
    if dev[-1].isdigit():
        answer = input("%s seems to be a partition. Are you sure you want to continue (y/Y)? " %(dev))
        check_answer(answer)

def show_drive_info(dev):
    try:
        subprocess.run(["hwinfo","--short","--disk","--only",dev])
    except FileNotFoundError:
        # hwinfo not installed
        print("hwinfo seems not installed")
    finally:
        answer = input(dev+" correct disk (y/Y)? ")
        check_answer(answer)

def change_drive(dev, hexval, offset_dec):
    try:
        with open(dev,"r+b") as file:
            for v in hexval:
                if isinstance(v,str):
                    # convert to integer
                    v = int(v,16)
                byteval = int.to_bytes(v,1,sys.byteorder)
                file.seek(offset_dec)
                file.write(byteval)
                print("wrote {0:X} to offset {1:X} (decimal {2})".format(v,offset_dec,offset_dec))
                offset_dec += 1
    except Exception as e:
        print("ERROR. something went wrong\n%s" %(str(e)))
        exit(2)


def main(args):
    check_root()
    check_dev(args.device)
    if args.skip_hwinfo is not True:
        show_drive_info(args.device)
    if args.hide and args.unhide:
        print("use either --hide OR --unhide NOT both")
        exit()
    if args.hide:
        change_drive(args.device,HIDE_VAL,STARTING_OFFSET)
        print("drive %s hidden" %args.device)
    elif args.unhide:
        change_drive(args.device,SHOW_VAL,STARTING_OFFSET)
        print("drive %s unhidden" %args.device)

if __name__ == "__main__":
    main(create_parser())