main, lsb и dct под все типы сообщений и картинок

core/lsb.py

@@ -1,132 +1,93 @@
 """
 Модуль LSB (Least Significant Bit) стеганографии.
-
-Скрывает текст в изображении путем замены последних битов пикселей.
-Каждый пиксель RGB хранит 3 бита информации (по одному в каждом канале).
+Поддержка цветных изображений и русских текстов.
 """
 
 import numpy as np
 from PIL import Image
 from .utils import text_to_bits, bits_to_text
 
-
 def encode_lsb(image_path: str, message: str, output_path: str) -> bool:
-    """
-    Скрывает текстовое сообщение в изображении методом LSB.
-
-    Алгоритм:
-    1. Преобразует текст в битовую строку
-    2. Добавляет стоп-маркер (8 нулевых битов) для обозначения конца
-    3. Проходит по всем пикселям изображения
-    4. В каждом канале (R,G,B) заменяет последний бит на бит сообщения
-
-    Аргументы:
-        image_path (str): Путь к исходному изображению
-        message (str): Текст для сокрытия
-        output_path (str): Путь для сохранения результата
-
-    Возвращает:
-        bool: True если успешно, False если сообщение слишком длинное
-    """
-    # Загружаем изображение
+    """Скрывает сообщение в изображении методом LSB."""
     img = Image.open(image_path).convert('RGB')
     pixels = np.array(img, dtype=np.uint8)
     height, width, channels = pixels.shape
 
-    # Преобразуем сообщение в биты и добавляем стоп-маркер
-    bits = text_to_bits(message) + '00000000'
+    # Формат: [длина сообщения: 32 бита] + [сообщение]
+    msg_bytes = message.encode('utf-8')
+    msg_length = len(msg_bytes)
+    length_bits = format(msg_length, '032b')
+    message_bits = text_to_bits(message)
+    all_bits = length_bits + message_bits
+    bit_list = [int(b) for b in all_bits]
+    total_bits = len(bit_list)
 
-    # Проверяем, поместится ли сообщение
+    # Проверка вместимости
     max_bits = height * width * channels
-    if len(bits) > max_bits:
+    if total_bits > max_bits:
         print(f"Ошибка: сообщение слишком большое.")
-        print(f"Доступно битов: {max_bits}, требуется: {len(bits)}")
+        print(f"Доступно: {max_bits} бит, требуется: {total_bits}")
         return False
 
-    # Прячем биты в изображение
+    # Внедрение битов
     bit_index = 0
-    total_bits = len(bits)
-
     for i in range(height):
         for j in range(width):
             for c in range(channels):
                 if bit_index >= total_bits:
                     break
-
                 pixel_value = pixels[i, j, c]
-                bit = int(bits[bit_index])
-
-                # Обнуляем последний бит и устанавливаем нужный
+                bit = bit_list[bit_index]
                 pixels[i, j, c] = (pixel_value & 0xFE) | bit
                 bit_index += 1
-
             if bit_index >= total_bits:
                 break
         if bit_index >= total_bits:
             break
 
-    # Сохраняем измененное изображение
     result_img = Image.fromarray(pixels, mode='RGB')
     result_img.save(output_path)
-
     print(f"Успешно! Спрятано {total_bits} бит ({total_bits // 8} байт)")
     return True
 
 
 def decode_lsb(image_path: str) -> str:
-    """
-    Извлекает скрытое сообщение из изображения методом LSB.
-
-    Алгоритм:
-    1. Проходит по всем пикселям изображения
-    2. Из каждого канала извлекает последний бит
-    3. Собирает биты в байты
-    4. Останавливается при обнаружении стоп-маркера (8 нулевых битов)
-
-    Аргументы:
-        image_path (str): Путь к изображению со скрытым сообщением
-
-    Возвращает:
-        str: Извлеченное сообщение
-    """
-    # Загружаем изображение
+    """Извлекает сообщение из изображения методом LSB."""
     img = Image.open(image_path).convert('RGB')
     pixels = np.array(img, dtype=np.uint8)
     height, width, channels = pixels.shape
 
-    # Извлекаем биты
-    extracted_bits = []
-    stop_counter = 0
-    stop_marker_length = 8
-
+    # Извлекаем все биты
+    all_bits = []
     for i in range(height):
         for j in range(width):
             for c in range(channels):
-                # Извлекаем последний бит
-                bit = pixels[i, j, c] & 1
-                extracted_bits.append(bit)
+                all_bits.append(pixels[i, j, c] & 1)
 
-                # Проверяем стоп-маркер
-                if bit == 0:
-                    stop_counter += 1
-                    if stop_counter == stop_marker_length:
-                        # Убираем стоп-маркер из результата
-                        extracted_bits = extracted_bits[:-stop_marker_length]
-                        break
-                else:
-                    stop_counter = 0
-
-            if stop_counter == stop_marker_length:
-                break
-        if stop_counter == stop_marker_length:
-            break
-
-    if len(extracted_bits) == 0:
+    # Читаем длину сообщения (первые 32 бита)
+    if len(all_bits) < 32:
         return ""
-
-    # Преобразуем биты в строку и декодируем
-    bits_string = ''.join(str(bit) for bit in extracted_bits)
-
+    
+    length_bits = all_bits[:32]
+    length_str = ''.join(str(b) for b in length_bits)
+    msg_length = int(length_str, 2)
+    
+    # Проверяем, что длина разумная
+    if msg_length > (len(all_bits) - 32) // 8:
+        return ""
+    
+    # Читаем сообщение
+    message_bits = all_bits[32:32 + msg_length * 8]
+    
+    # Дополняем до кратности 8 (на всякий случай)
+    remainder = len(message_bits) % 8
+    if remainder != 0:
+        message_bits.extend([0] * (8 - remainder))
+    
+    if len(message_bits) == 0:
+        return ""
+    
+    bits_string = ''.join(str(b) for b in message_bits)
     try:
         return bits_to_text(bits_string)
     except Exception as e: