Сделал CLI интерфейс для LSB
encode, decode, capacity, psnr;
This commit is contained in:
Игорь
211
main.py
211
main.py
@@ -0,0 +1,211 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
CLI интерфейс для стеганографических методов LSB и DCT.
|
||||
|
||||
Использование:
|
||||
python main.py encode <method> <image> <output> [-m MESSAGE] [-f FILE]
|
||||
python main.py decode <method> <image> [-o OUTPUT]
|
||||
|
||||
Примеры:
|
||||
python main.py encode lsb image.png result.png -m "Secret message"
|
||||
python main.py encode lsb image.png result.png -f secret.txt
|
||||
python main.py decode lsb result.png
|
||||
python main.py decode lsb result.png -o extracted.txt
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
import os
|
||||
from core.lsb import encode_lsb, decode_lsb
|
||||
from core.utils import calculate_capacity, calculate_psnr
|
||||
|
||||
|
||||
def read_message_from_file(file_path: str) -> str:
|
||||
"""Читает сообщение из текстового файла."""
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
return f.read()
|
||||
|
||||
|
||||
def save_message_to_file(file_path: str, message: str) -> None:
|
||||
"""Сохраняет извлеченное сообщение в файл."""
|
||||
with open(file_path, 'w', encoding='utf-8') as f:
|
||||
f.write(message)
|
||||
print(f"Сообщение сохранено в: {file_path}")
|
||||
|
||||
|
||||
def show_capacity(image_path: str) -> None:
|
||||
"""Показывает максимальную вместимость изображения для LSB метода."""
|
||||
try:
|
||||
capacity = calculate_capacity(image_path)
|
||||
print(f"Максимальная вместимость изображения: {capacity} байт")
|
||||
print(f"Это примерно {capacity / 1024:.2f} КБ текста")
|
||||
except Exception as e:
|
||||
print(f"Ошибка при расчете вместимости: {e}")
|
||||
|
||||
|
||||
def show_psnr(original_path: str, modified_path: str) -> None:
|
||||
"""Показывает PSNR между двумя изображениями."""
|
||||
try:
|
||||
psnr = calculate_psnr(original_path, modified_path)
|
||||
print(f"PSNR: {psnr:.2f} dB")
|
||||
|
||||
if psnr > 40:
|
||||
print("Качество: Отличное (изменения практически незаметны)")
|
||||
elif psnr > 30:
|
||||
print("Качество: Хорошее")
|
||||
elif psnr > 20:
|
||||
print("Качество: Заметные искажения")
|
||||
else:
|
||||
print("Качество: Сильные искажения")
|
||||
except Exception as e:
|
||||
print(f"Ошибка при расчете PSNR: {e}")
|
||||
|
||||
|
||||
def encode_command(args) -> None:
|
||||
"""Обрабатывает команду encode."""
|
||||
# Получаем сообщение
|
||||
if args.message and args.file:
|
||||
print("Ошибка: используйте либо -m, либо -f, но не оба")
|
||||
sys.exit(1)
|
||||
|
||||
if args.message:
|
||||
message = args.message
|
||||
elif args.file:
|
||||
try:
|
||||
message = read_message_from_file(args.file)
|
||||
except FileNotFoundError:
|
||||
print(f"Ошибка: файл не найден - {args.file}")
|
||||
sys.exit(1)
|
||||
else:
|
||||
print("Ошибка: укажите сообщение через -m или -f")
|
||||
sys.exit(1)
|
||||
|
||||
# Проверяем, поместится ли сообщение (только для LSB)
|
||||
if args.method == 'lsb':
|
||||
try:
|
||||
capacity = calculate_capacity(args.image)
|
||||
msg_bytes = len(message.encode('utf-8'))
|
||||
if msg_bytes > capacity:
|
||||
print(f"Предупреждение: сообщение может не поместиться!")
|
||||
print(f"Доступно: {capacity} байт, требуется: {msg_bytes} байт")
|
||||
response = input("Продолжить? (y/n): ")
|
||||
if response.lower() != 'y':
|
||||
print("Отменено")
|
||||
sys.exit(0)
|
||||
except Exception as e:
|
||||
print(f"Ошибка при проверке вместимости: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
# Выполняем кодирование
|
||||
print(f"Кодирование методом {args.method.upper()}...")
|
||||
|
||||
if args.method == 'lsb':
|
||||
success = encode_lsb(args.image, message, args.output)
|
||||
elif args.method == 'dct':
|
||||
print("Метод DCT пока не реализован")
|
||||
success = False
|
||||
else:
|
||||
print(f"Неизвестный метод: {args.method}")
|
||||
success = False
|
||||
|
||||
if success:
|
||||
print("Кодирование завершено успешно!")
|
||||
|
||||
if args.psnr:
|
||||
show_psnr(args.image, args.output)
|
||||
else:
|
||||
print("Ошибка при кодировании")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def decode_command(args) -> None:
|
||||
"""Обрабатывает команду decode."""
|
||||
print(f"Декодирование методом {args.method.upper()}...")
|
||||
|
||||
if args.method == 'lsb':
|
||||
message = decode_lsb(args.image)
|
||||
elif args.method == 'dct':
|
||||
print("Метод DCT пока не реализован")
|
||||
message = None
|
||||
else:
|
||||
print(f"Неизвестный метод: {args.method}")
|
||||
sys.exit(1)
|
||||
|
||||
if message:
|
||||
if args.output:
|
||||
save_message_to_file(args.output, message)
|
||||
else:
|
||||
print("\n" + "=" * 50)
|
||||
print("Извлеченное сообщение:")
|
||||
print("=" * 50)
|
||||
print(message)
|
||||
print("=" * 50)
|
||||
else:
|
||||
print("Не удалось извлечь сообщение или сообщение пустое")
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Главная функция: парсит аргументы и запускает команду."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Steganography tool: hide data in images using LSB and DCT",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog="""
|
||||
Примеры:
|
||||
%(prog)s encode lsb image.png result.png -m "Secret message"
|
||||
%(prog)s encode lsb image.png result.png -f secret.txt
|
||||
%(prog)s decode lsb result.png
|
||||
%(prog)s decode lsb result.png -o extracted.txt
|
||||
%(prog)s capacity image.png
|
||||
%(prog)s psnr original.png modified.png
|
||||
"""
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(dest='command', help='Команды')
|
||||
|
||||
# Команда encode
|
||||
encode_parser = subparsers.add_parser('encode', help='Скрыть сообщение')
|
||||
encode_parser.add_argument('method', choices=['lsb', 'dct'],
|
||||
help='Метод стеганографии')
|
||||
encode_parser.add_argument('image', help='Путь к исходному изображению')
|
||||
encode_parser.add_argument('output', help='Путь для сохранения результата')
|
||||
encode_parser.add_argument('-m', '--message', help='Текст сообщения')
|
||||
encode_parser.add_argument('-f', '--file', help='Файл с сообщением')
|
||||
encode_parser.add_argument('--psnr', action='store_true',
|
||||
help='Показать PSNR после кодирования')
|
||||
|
||||
# Команда decode
|
||||
decode_parser = subparsers.add_parser('decode', help='Извлечь сообщение')
|
||||
decode_parser.add_argument('method', choices=['lsb', 'dct'],
|
||||
help='Метод стеганографии')
|
||||
decode_parser.add_argument('image', help='Путь к изображению')
|
||||
decode_parser.add_argument('-o', '--output', help='Сохранить в файл')
|
||||
|
||||
# Команда capacity
|
||||
capacity_parser = subparsers.add_parser('capacity',
|
||||
help='Показать вместимость')
|
||||
capacity_parser.add_argument('image', help='Путь к изображению')
|
||||
|
||||
# Команда psnr
|
||||
psnr_parser = subparsers.add_parser('psnr',
|
||||
help='Сравнить изображения')
|
||||
psnr_parser.add_argument('original', help='Исходное изображение')
|
||||
psnr_parser.add_argument('modified', help='Измененное изображение')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.command == 'encode':
|
||||
encode_command(args)
|
||||
elif args.command == 'decode':
|
||||
decode_command(args)
|
||||
elif args.command == 'capacity':
|
||||
show_capacity(args.image)
|
||||
elif args.command == 'psnr':
|
||||
show_psnr(args.original, args.modified)
|
||||
else:
|
||||
parser.print_help()
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user