input: cover_image, payload
payload_bits = header(len(payload)) + to_bits(payload)
assert len(payload_bits) <= capacity(cover_image)
indices = prng_indices(key) or range(capacity)
for i, idx in enumerate(indices):
    if i == len(payload_bits): break
    channel_value = get_channel(cover_image, idx)
    set_channel(cover_image, idx, (channel_value & ~1) | payload_bits[i])
save(cover_image, 'stego.png')

from PIL import Image

def text_to_bits(text):
    """将文本转换为位串"""
    bits = []
    for char in text:
        binary = bin(ord(char))[2:].zfill(8)  # 将字符转换为8位二进制
        bits.extend([int(bit) for bit in binary])
    return bits

def hide_data(image_path, message, output_path):
    """在图像中隐藏信息"""
    # 打开图像
    img = Image.open(image_path)
    width, height = img.size
    img_data = list(img.getdata())

    # 转换消息为位
    binary_message = text_to_bits(message)

    # 在消息前添加长度信息（32位）
    message_length = len(binary_message)
    length_bits = [int(bit) for bit in bin(message_length)[2:].zfill(32)]

    binary_data = length_bits + binary_message

    if len(binary_data) > len(img_data) * 3:  # 每个像素3个通道
        raise ValueError("消息太长，无法在图像中隐藏")

    # 修改像素的LSB
    new_img_data = []
    data_index = 0

    for pixel in img_data:
        # 对RGB通道分别处理
        new_pixel = list(pixel)

        for i in range(3):  # RGB三个通道
            if data_index < len(binary_data):
                # 修改最低有效位
                new_pixel[i] = (new_pixel[i] & ~1) | binary_data[data_index]
                data_index += 1
            else:
                break

        new_img_data.append(tuple(new_pixel))

        if data_index >= len(binary_data):
            # 所有数据已嵌入，添加剩余像素
            new_img_data.extend(img_data[len(new_img_data):])
            break

    # 创建新图像并保存
    new_img = Image.new(img.mode, img.size)
    new_img.putdata(new_img_data)
    new_img.save(output_path)

    return f"消息已成功隐藏在图像中。总共使用了{data_index}位。"

def extract_data(image_path):
    """从图像中提取隐藏的信息"""
    # 打开图像
    img = Image.open(image_path)
    img_data = list(img.getdata())

    # 提取所有LSB
    binary_data = []
    for pixel in img_data:
        for i in range(3):  # RGB三个通道
            binary_data.append(pixel[i] & 1)

    # 提取消息长度
    length_bits = binary_data[:32]
    message_length = int(''.join(map(str, length_bits)), 2)

    # 提取消息
    message_bits = binary_data[32:32 + message_length]

    # 将位转换回字符
    message = ""
    for i in range(0, len(message_bits), 8):
        if i + 8 <= len(message_bits):
            byte = message_bits[i:i+8]
            char = chr(int(''.join(map(str, byte)), 2))
            message += char

    return message

# 隐藏信息
result = hide_data("original_image.png", "这是一条隐藏消息", "stego_image.png")
print(result)

# 提取信息
extracted_message = extract_data("stego_image.png")
print("提取的消息:", extracted_message)