diff --git a/main.py b/main.py index e69de29..d43696f 100644 --- a/main.py +++ b/main.py @@ -0,0 +1,211 @@ +#!/usr/bin/env python3 +""" +CLI интерфейс для стеганографических методов LSB и DCT. + +Использование: + python main.py encode [-m MESSAGE] [-f FILE] + python main.py decode [-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() \ No newline at end of file