from PIL import Image

def lsb_encode(image_path, secret_text, output_path):
    """
    将 secret_text 编码到图片的 LSB 中
    数据格式: [32bit 长度] + [utf-8 编码的文本]
    """
    img = Image.open(image_path).convert('RGB')
    pixels = list(img.getdata())

    # 将文本转为 bit 流
    text_bytes = secret_text.encode('utf-8')
    length_bytes = len(text_bytes).to_bytes(4, 'big')
    data = length_bytes + text_bytes
    bit_stream = ''.join(f'{byte:08b}' for byte in data)

    if len(bit_stream) > len(pixels) * 3:
        raise ValueError('图片容量不足，需要更大的图片')

    new_pixels = []
    bit_idx = 0

    for r, g, b in pixels:
        if bit_idx < len(bit_stream):
            r = (r & 0xFE) | int(bit_stream[bit_idx])
            bit_idx += 1
        if bit_idx < len(bit_stream):
            g = (g & 0xFE) | int(bit_stream[bit_idx])
            bit_idx += 1
        if bit_idx < len(bit_stream):
            b = (b & 0xFE) | int(bit_stream[bit_idx])
            bit_idx += 1
        new_pixels.append((r, g, b))

    img.putdata(new_pixels)
    img.save(output_path, 'PNG')  # 必须用无损格式
    print(f'已编码 {len(text_bytes)} 字节 → {output_path}')

# 使用示例
lsb_encode('original.png', 'Hello, 隐写!', 'stego.png')

def lsb_decode(stego_path):
    """从 LSB 隐写图片中提取隐藏文本"""
    img = Image.open(stego_path).convert('RGB')
    pixels = list(img.getdata())

    # 提取所有 LSB
    bits = []
    for r, g, b in pixels:
        bits.append(str(r & 1))
        bits.append(str(g & 1))
        bits.append(str(b & 1))

    bit_str = ''.join(bits)

    # 先读取 4 字节长度头
    length_bits = bit_str[:32]
    text_length = int(length_bits, 2)

    # 读取正文
    text_bits = bit_str[32:32 + text_length * 8]
    text_bytes = bytes(
        int(text_bits[i:i+8], 2) for i in range(0, len(text_bits), 8)
    )

    return text_bytes.decode('utf-8')

# 使用示例
secret = lsb_decode('stego.png')
print(f'提取到: {secret}')

def embed_image(cover_path, secret_path, output_path):
    cover = Image.open(cover_path).convert('RGB')
    secret = Image.open(secret_path).convert('RGB').resize(cover.size)

    cover_pixels = list(cover.getdata())
    secret_pixels = list(secret.getdata())

    new_pixels = []
    for (cr, cg, cb), (sr, sg, sb) in zip(cover_pixels, secret_pixels):
        nr = (cr & 0xF0) | (sr >> 4)   # 封面高4bit + 秘密高4bit
        ng = (cg & 0xF0) | (sg >> 4)
        nb = (cb & 0xF0) | (sb >> 4)
        new_pixels.append((nr, ng, nb))

    img = Image.new('RGB', cover.size)
    img.putdata(new_pixels)
    img.save(output_path, 'PNG')


def extract_image(stego_path, output_path):
    img = Image.open(stego_path).convert('RGB')
    pixels = list(img.getdata())

    new_pixels = []
    for r, g, b in pixels:
        sr = (r & 0x0F) << 4  # 提取低4bit并左移
        sg = (g & 0x0F) << 4
        sb = (b & 0x0F) << 4
        new_pixels.append((sr, sg, sb))

    result = Image.new('RGB', img.size)
    result.putdata(new_pixels)
    result.save(output_path, 'PNG')

import struct
import zlib

def fix_png_size(filename):
    with open(filename, 'rb') as f:
        data = f.read()

    # 读取 IHDR 中的宽高
    ihdr_start = data.find(b'IHDR')
    width = struct.unpack('>I', data[ihdr_start+4:ihdr_start+8])[0]
    height = struct.unpack('>I', data[ihdr_start+8:ihdr_start+12])[0]

    # 暴力枚举正确的宽高
    for w in range(width, width + 2000):
        for h in range(height, height + 2000):
            # 用新的宽高计算 CRC
            new_ihdr = struct.pack('>IIBBBBB', w, h, *data[ihdr_start+12:ihdr_start+17])
            if zlib.crc32(b'IHDR' + new_ihdr) == struct.unpack('>I', data[ihdr_start+17:ihdr_start+21])[0]:
                return w, h
    return None