from ida_segment import *
from ida_bytes import *
from idc import *


def setup_wifi_regions():
    WIFI_BASE = 0x40080000
    SIE_BASE = 0x400C0000

    WIFI_SIZE = SIE_BASE - WIFI_BASE
    SIE_SIZE = 0x1000

    print("Setting up WiFi memory regions...")

    if not add_segm(0, WIFI_BASE, SIE_BASE, "WIFI", "DATA"):
        print("Failed to create WIFI segment")
        return

    if not add_segm(0, SIE_BASE, SIE_BASE + SIE_SIZE, "SIE", "DATA"):
        print("Failed to create SIE segment")
        return

    print(f"Memory regions created:")
    print(f"WIFI: {hex(WIFI_BASE)} - {hex(SIE_BASE-1)}")
    print(f"SIE:  {hex(SIE_BASE)} - {hex(SIE_BASE + SIE_SIZE-1)}")


if __name__ == "__main__":
    setup_wifi_regions()