#!/usr/bin/env python3
import sys
import os
import random
import argparse
import re

UGREEN_BIN = "ugreen.bin"
A1277_BIN = "a1277.bin"

OUIS = {
    "nintendo": [[0x00, 0x1E, 0xA9], [0x00, 0x22, 0xAA], [0x00, 0x24, 0x44], [0x00, 0x09, 0xBF]],
    "apple": [[0x00, 0x0A, 0x27], [0x00, 0x14, 0x51], [0x00, 0x23, 0xDF], [0x04, 0x0C, 0xCE]],
    "asix": [[0x00, 0x0E, 0xC6], [0x00, 0x13, 0x3B]]
}

def generate_mac(vendor_key):
    oui_list = OUIS.get(vendor_key, OUIS["asix"])
    selected_oui = random.choice(oui_list)
    return "".join(f"{b:02X}" for b in selected_oui + [random.randint(0x00, 0xFF) for _ in range(3)])

def extract_a1277_data(filepath):
    """Strictly extracts MAC and Serial from a genuine Apple EEPROM dump."""
    with open(filepath, "rb") as f:
        data = bytearray(f.read())
        
    if len(data) != 512:
        sys.exit(f"Error: {filepath} is not exactly 512 bytes. Extraction aborted.")

    # Extract MAC
    raw_mac = data[8:14]
    unswapped_mac = bytearray(6)
    for i in range(0, 6, 2):
        unswapped_mac[i] = raw_mac[i+1]
        unswapped_mac[i+1] = raw_mac[i]
    mac = "".join(f"{b:02X}" for b in unswapped_mac)

    # Extract Serial
    serial_len = data[0x14]
    serial_offset = data[0x15] * 2
    
    if serial_offset == 0 or serial_offset + serial_len > 512:
        sys.exit("Error: Corrupt Serial Pointer detected in A1277 binary.")
        
    if data[serial_offset] != 0x03 or data[serial_offset+1] != serial_len:
        sys.exit("Error: Invalid String Descriptor Header for Serial Number.")

    raw_serial = data[serial_offset+2 : serial_offset + serial_len]
    unswapped_serial = bytearray(len(raw_serial))
    for i in range(0, len(raw_serial), 2):
        unswapped_serial[i] = raw_serial[i+1]
        unswapped_serial[i+1] = raw_serial[i]
        
    try:
        serial = unswapped_serial.decode('utf-16le').strip('\x00')
    except UnicodeDecodeError:
        sys.exit("Error: Extracted Serial number is not valid UTF-16LE data.")

    return mac, serial

def validate_hex(value, length, name):
    """Strictly validates hexadecimal identifiers."""
    clean_val = re.sub(r'[^A-Fa-f0-9]', '', value).upper()
    if len(clean_val) != length:
        sys.exit(f"Error: {name} must be exactly {length} valid hex characters. Received '{value}'.")
    return clean_val

def validate_string_length(value, max_length, name):
    """Enforces EEPROM boundary limits on UTF-16 strings."""
    if len(value) > max_length:
        sys.exit(f"Error: {name} exceeds the maximum hardware limit of {max_length} characters.")
    return value

def main():
    parser = argparse.ArgumentParser(description="AX88772 Hardware EEPROM Patcher")
    parser.add_argument("-b", "--base", choices=["nintendo", "apple", "ax88772", "ax88772a"], 
                        help="Baseline preset to populate default identifiers.")
    parser.add_argument("-v", "--vid", type=str, help="Override Vendor ID (4 hex chars)")
    parser.add_argument("-d", "--pid", type=str, help="Override Product ID (4 hex chars)")
    parser.add_argument("-a", "--mac", type=str, help="Override MAC Address (12 hex chars)")
    parser.add_argument("-m", "--manufacturer", type=str, help="Override Manufacturer String")
    parser.add_argument("-p", "--product", type=str, help="Override Product String")
    parser.add_argument("-s", "--serial", type=str, help="Override Serial Number")
    
    args = parser.parse_args()

    # 1. HARDWARE FOUNDATION CHECK
    if not os.path.exists(UGREEN_BIN):
        sys.exit(f"FATAL ERROR: Base reference file '{UGREEN_BIN}' not found in the current directory.\n"
                 f"The script requires the unlocked UGREEN firmware as a structural canvas to bypass signature checks.")

    # 2. ESTABLISH BASELINE DEFAULTS
    preset = args.base if args.base else "ax88772"
    oui_key = "asix"
    
    vid = "0B95"
    pid = "7720"
    mfg = "ASIX Elec. Corp."
    prod = "AX88772"

    if preset == "nintendo":
        oui_key = "nintendo"
        prod = "AX88772 "
    elif preset == "apple":
        oui_key = "apple"
        vid = "05AC"
        pid = "1402"
        mfg = "Apple Inc."
        prod = "Apple USB Ethernet Adapter"
    elif preset == "ax88772a":
        pid = "772A"
        prod = "AX88772A"

    mac = generate_mac(oui_key)
    serial = mac[-6:]

    # 3. A1277 GRAFT EXTRACTION OR NOTICE
    if os.path.exists(A1277_BIN):
        print(f"[*] Notice: Found {A1277_BIN}. Extracting genuine Apple MAC and Serial...")
        extracted_mac, extracted_serial = extract_a1277_data(A1277_BIN)
        mac = extracted_mac
        serial = extracted_serial
    else:
        print(f"[*] Notice: {A1277_BIN} not found. Proceeding with generated {oui_key.upper()} OUI MAC and derived serial.")

    # 4. CLI OVERRIDES (Highest Priority)
    if args.vid: vid = args.vid
    if args.pid: pid = args.pid
    if args.mac: mac = args.mac
    if args.manufacturer: mfg = args.manufacturer
    if args.product: prod = args.product
    if args.serial: serial = args.serial

    # 5. STRICT VALIDATION
    vid = validate_hex(vid, 4, "Vendor ID (VID)")
    pid = validate_hex(pid, 4, "Product ID (PID)")
    mac = validate_hex(mac, 12, "MAC Address")
    
    mfg = validate_string_length(mfg, 31, "Manufacturer String")
    prod = validate_string_length(prod, 31, "Product String")
    serial = validate_string_length(serial, 31, "Serial Number")

    # 6. PHYSICAL INJECTION
    with open(UGREEN_BIN, "rb") as f:
        eeprom = bytearray(f.read())

    # VID/PID
    vid_bytes = bytes.fromhex(vid)
    pid_bytes = bytes.fromhex(pid)
    eeprom[0x48:0x4A] = vid_bytes
    eeprom[0x4A:0x4C] = pid_bytes
    eeprom[0x88:0x8A] = vid_bytes
    eeprom[0x8A:0x8C] = pid_bytes

    # MAC
    mac_bytes = bytes.fromhex(mac)
    swapped_mac = bytearray(6)
    for i in range(0, 6, 2):
        swapped_mac[i] = mac_bytes[i+1]
        swapped_mac[i+1] = mac_bytes[i]
    eeprom[8:14] = swapped_mac

    # Serial
    serial_raw = bytearray(serial.encode('utf-16le'))
    for i in range(0, len(serial_raw), 2):
        serial_raw[i], serial_raw[i+1] = serial_raw[i+1], serial_raw[i]
    s_len = len(serial_raw) + 2
    eeprom[0x32] = 0x03
    eeprom[0x33] = s_len
    eeprom[0x34:0x34+len(serial_raw)] = serial_raw
    for i in range(0x34 + len(serial_raw), 0x40): eeprom[i] = 0x00

    # Clean Vanity Area
    for i in range(0xC0, 0x140): eeprom[i] = 0xFF

    # Manufacturer
    m_raw = bytearray(mfg.encode('utf-16le'))
    for i in range(0, len(m_raw), 2):
        m_raw[i], m_raw[i+1] = m_raw[i+1], m_raw[i]
    m_len = len(m_raw) + 2
    m_offset = 0xC0
    eeprom[m_offset] = 0x03
    eeprom[m_offset+1] = m_len
    eeprom[m_offset+2:m_offset+2+len(m_raw)] = m_raw

    # Product
    p_raw = bytearray(prod.encode('utf-16le'))
    for i in range(0, len(p_raw), 2):
        p_raw[i], p_raw[i+1] = p_raw[i+1], p_raw[i]
    p_len = len(p_raw) + 2
    p_offset = m_offset + m_len
    eeprom[p_offset] = 0x03
    eeprom[p_offset+1] = p_len
    eeprom[p_offset+2:p_offset+2+len(p_raw)] = p_raw

    # Pointer Table & Word Count
    eeprom[0x10] = m_len
    eeprom[0x11] = m_offset // 2
    eeprom[0x12] = p_len
    eeprom[0x13] = p_offset // 2
    eeprom[0x14] = s_len
    eeprom[0x15] = 0x32 // 2

    required_words = (p_offset + p_len) // 2
    if required_words > eeprom[0]:
        eeprom[0] = required_words

    # 7. WRITE OUTPUT
    safe_mac_str = "-".join(mac[i:i+2] for i in range(0, 12, 2))
    target_filename = f"patched_ax88772_{safe_mac_str}.bin"
    
    with open(target_filename, "wb") as f:
        f.write(eeprom)
        
    print(f"[+] Success! Configuration physically patched and saved to: {target_filename}")
    print(f"    MAC: {safe_mac_str} | VID: {vid} | PID: {pid}")

if __name__ == "__main__":
    main()
