x230patcher/clishared.py

from pathlib import Path
import filecmp

from romtools import *

def cli_join_rom(clean_path):
    print("Joining dumps in to full firmware ROM... ", end="")
    with open(clean_path / "bottom.bin", "rb") as bottom, open(clean_path / "top.bin", "rb") as top, open(clean_path / "clean.rom", "wb+") as rom:
        rom.write(bottom.read() + top.read())
    print("Done.")
    
def cli_split_rom(patched_path):
    print("Splitting full rom in to images for each chip... ", end="")
    with open(patched_path / "bottom.bin", "wb+") as bottom, open(patched_path / "top.bin", "wb+") as top, open(patched_path / "patched.rom", "rb") as rom:
        rom_array = rom.read()
        bottom.write(rom_array[:0x800000])
        top.write(rom_array[0x800000:])
    print("Done.")

def cli_connect_rom(chip_loc):
    input(f"Please connect programming clip to {chip_loc} flash chip. Press Enter when done...")
    cli_reconnect_rom(chip_loc)
    
def cli_reconnect_rom(chip_loc):
    print("Checking if ROM is connected... ", end="")
    try:
        check_rom(chip_loc)
        print("ROM is reachable.")
    except:
        response = input("ROM is not reachable. Try again? (Y/q): ")
        if response in ['q', 'Q']:
            raise ROMError("Could not connect to ROM. User elected to quit.")
        else:
            cli_reconnect_rom(chip_loc)
        
def cli_read_rom(chip_loc, clean_path, verify=True):
    files = [chip_loc + ".bin", chip_loc + "_2.bin"] if verify else [chip_loc + ".bin"]
    for count, file_name in enumerate(files):
        print(f"Reading {chip_loc} ROM, pass {count+1} of {len(files)}... ", end="", flush=True)
        read_rom(chip_loc, clean_path, file_name)
        print("Success.")
            
    if verify: 
        print("Comparing ROM files... ", end="")
        if filecmp.cmp(clean_path / files[0], clean_path / files[1]):
            print("Files match.")
        else:
            raise ROMError("ROM files do not match. Read the log and try again.")
    else:
        print("ROM file will not be verified, proceed at your own risk!")
        
    print(f"{chip_loc.capitalize()} ROM read successfully.")
    
def cli_write_rom(chip_loc, patched_path, verify=True):
    file_name = chip_loc + '_patched.bin'
    print(f"Writing {chip_loc} ROM (takes a while)... ", end="", flush=True)
    write_rom(chip_loc, patched_path, file_name)
    print("Success.")
        
    print(f"{chip_loc.capitalize()} ROM written successfully.")

def cli_choose_patches():
    which_patches = input("Patch (W)LAN whitelist, (A)dvanced menu or (B)oth? ")
    
    if which_patches.strip()[0] in ['w', 'W']:
        patches = ["wifi"]
    elif which_patches.strip()[0] in ['a', 'A']:
        patches = ["advanced"]
    elif which_patches.strip()[0] in ['b', 'B']:
        patches = ["wifi", "advanced"]
    else:
        raise ValueError("Selection invalid.")
    
    return patches
    
def cli_patch_rom(patches, clean_path, patched_path):
    patch_rom(patches, clean_path / "top.bin", patched_path / "top_patched.bin")
    print("ROM patched successfully.")