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

core/dct.py

@@ -1,5 +1,6 @@
 """
 Модуль DCT стеганографии с квантованием (JPEG-style)
+Оптимизирован для длинных сообщений
 """
 
 import numpy as np
@@ -9,16 +10,16 @@ from .utils import text_to_bits, bits_to_text
 
 BLOCK_SIZE = 8
 
-# Стандартная таблица квантования JPEG для качества 50%
+# Более мягкая таблица квантования (качество ~70-80%)
 QUANT_TABLE = np.array([
-    [16, 11, 10, 16, 24, 40, 51, 61],
-    [12, 12, 14, 19, 26, 58, 60, 55],
-    [14, 13, 16, 24, 40, 57, 69, 56],
-    [14, 17, 22, 29, 51, 87, 80, 62],
-    [18, 22, 37, 56, 68, 109, 103, 77],
-    [24, 35, 55, 64, 81, 104, 113, 92],
-    [49, 64, 78, 87, 103, 121, 120, 101],
-    [72, 92, 95, 98, 112, 100, 103, 99]
+    [8, 6, 5, 8, 12, 20, 26, 31],
+    [6, 6, 7, 10, 13, 29, 30, 28],
+    [7, 7, 8, 12, 20, 29, 35, 28],
+    [7, 9, 11, 15, 26, 44, 40, 31],
+    [9, 11, 19, 28, 34, 55, 52, 39],
+    [12, 18, 28, 32, 41, 52, 57, 46],
+    [25, 32, 39, 44, 52, 61, 60, 51],
+    [36, 46, 48, 49, 56, 50, 52, 50]
 ])
 
 
@@ -40,7 +41,8 @@ def _get_zigzag_order() -> list:
 
 
 ZIGZAG_ORDER = _get_zigzag_order()
-MID_FREQ_INDICES = ZIGZAG_ORDER[1:40]
+# Используем больше коэффициентов (1-50)
+MID_FREQ_INDICES = ZIGZAG_ORDER[1:50]
 
 
 def _dct_quantize(block: np.ndarray) -> np.ndarray:
@@ -93,28 +95,21 @@ def _embed_bits_in_block(block_quant: np.ndarray, bits: list, bit_index: int) ->
 
 def encode_dct(image_path: str, message: str, output_path: str) -> bool:
     """Скрывает сообщение в изображении методом DCT."""
-    img = Image.open(image_path).convert('L')
+    img = Image.open(image_path).convert('RGB')
     pixels = np.array(img, dtype=np.float64)
-    height, width = pixels.shape
+    height, width, channels = pixels.shape
 
-    # Кодируем сообщение в байты
-    message_bytes = message.encode('utf-8')
-    msg_length = len(message_bytes)
-    
-    # Формат: [длина: 32 бита] + [сообщение]
-    # 32 бита = 4 байта, достаточно для сообщений до 4 ГБ
+    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)
 
-    # Проверяем вместимость
     blocks_per_row = width // BLOCK_SIZE
     blocks_per_col = height // BLOCK_SIZE
-    max_bits = blocks_per_row * blocks_per_col * len(MID_FREQ_INDICES)
+    max_bits = blocks_per_row * blocks_per_col * len(MID_FREQ_INDICES) * 3
 
     if total_bits > max_bits:
         print(f"Ошибка: сообщение слишком большое.")
@@ -129,16 +124,20 @@ def encode_dct(image_path: str, message: str, output_path: str) -> bool:
             if bit_index >= total_bits:
                 break
             
-            block = pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE]
-            quant_block = _dct_quantize(block)
-            quant_block, bit_index = _embed_bits_in_block(quant_block, bit_list, bit_index)
-            new_block = _idct_dequantize(quant_block)
-            modified_pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE] = new_block
+            for c in range(3):
+                if bit_index >= total_bits:
+                    break
+                    
+                block = pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE, c]
+                quant_block = _dct_quantize(block)
+                quant_block, bit_index = _embed_bits_in_block(quant_block, bit_list, bit_index)
+                new_channel = _idct_dequantize(quant_block)
+                modified_pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE, c] = new_channel
             
         if bit_index >= total_bits:
             break
 
-    result_img = Image.fromarray(modified_pixels.astype(np.uint8), mode='L')
+    result_img = Image.fromarray(modified_pixels.astype(np.uint8), mode='RGB')
     result_img.save(output_path)
     print(f"Успешно! Спрятано {total_bits} бит ({total_bits // 8} байт)")
     return True
@@ -146,21 +145,20 @@ def encode_dct(image_path: str, message: str, output_path: str) -> bool:
 
 def decode_dct(image_path: str) -> str:
     """Извлекает сообщение из изображения методом DCT."""
-    img = Image.open(image_path).convert('L')
+    img = Image.open(image_path).convert('RGB')
     pixels = np.array(img, dtype=np.float64)
-    height, width = pixels.shape
+    height, width, channels = pixels.shape
 
-    # Извлекаем все биты
     all_bits = []
     
     for i in range(0, height - BLOCK_SIZE + 1, BLOCK_SIZE):
         for j in range(0, width - BLOCK_SIZE + 1, BLOCK_SIZE):
-            block = pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE]
-            quant_block = _dct_quantize(block)
-            bits_from_block = _extract_bits_from_block(quant_block, len(MID_FREQ_INDICES))
-            all_bits.extend(bits_from_block)
+            for c in range(3):
+                channel_block = pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE, c]
+                quant_block = _dct_quantize(channel_block)
+                bits_from_block = _extract_bits_from_block(quant_block, len(MID_FREQ_INDICES))
+                all_bits.extend(bits_from_block)
 
-    # Читаем длину сообщения (первые 32 бита)
     if len(all_bits) < 32:
         return ""
     
@@ -168,14 +166,11 @@ def decode_dct(image_path: str) -> str:
     length_str = ''.join(str(b) for b in length_bits)
     msg_length = int(length_str, 2)
     
-    # Проверяем, что длина разумная
-    if msg_length > len(all_bits) // 8:
+    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))

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:

main.py

@@ -81,7 +81,7 @@ def encode_command(args) -> None:
         print("Ошибка: укажите сообщение через -m или -f")
         sys.exit(1)
 
-    # Проверяем, поместится ли сообщение (только для LSB)
+    # Проверяем вместимость (только для LSB)
     if args.method == 'lsb':
         try:
             capacity = calculate_capacity(args.image)
@@ -104,14 +104,12 @@ def encode_command(args) -> None:
         success = encode_lsb(args.image, message, args.output)
     elif args.method == 'dct':
         success = encode_dct(args.image, message, args.output)
-        success = False
     else:
         print(f"Неизвестный метод: {args.method}")
         success = False
 
     if success:
         print("Кодирование завершено успешно!")
-
         if args.psnr:
             show_psnr(args.image, args.output)
     else:
@@ -127,7 +125,6 @@ def decode_command(args) -> None:
         message = decode_lsb(args.image)
     elif args.method == 'dct':
         message = decode_dct(args.image)
-        message = None
     else:
         print(f"Неизвестный метод: {args.method}")
         sys.exit(1)