from stegano import lsb

class auth():
    def __init__(self):
        global username
        username = input('Enter USERNAME: ')

        if len(username.split()) > 0:
            main()
        else:
            exit()


class main():
    def __init__(self):
        menu = '''
        [1] - спрятать сообщение в изображении
        [2] - показать спрятанное сообщение
        [3] - спрятать сообщение в изображении (с подписью)
        '''
        print(menu)

        cmd = input('CMD: ')

        if cmd == '1':
            image = input('Введите абсолютный путь до файла PNG or JPG: ')
            msg = input('Введите секретное сообщение: ')

            if len(image.split()) > 0 and len(msg.split()) > 0:
                sc_img = lsb.hide(image, msg)
                sc_img.save('secret_image.png')
                print('Сообщение было спрятано в secret_image.png')

            else:
                exit()

        elif cmd == '2':
            image = input('Введите абсолютный путь до файла PNG or JPG: ')

            if len(image.split()) > 0:
                sc_message = lsb.reveal(image)
                print(f'Секретное сообщение:\n{sc_message}')

            else:
                exit()

        elif cmd == '3':
            image = input('Введите абсолютный путь до файла PNG or JPG: ')
            msg = input('Введите секретное сообщение: ')

            if len(image.split()) > 0 and len(msg.split()) > 0:
                sc_img = lsb.hide(image, message=f'{msg}\n\nby {username}')
                sc_img.save('secret_image.png')
                print('Сообщение было спрятано в secret_image.png')

            else:
                exit()


if __name__ == '__main__':
    auth()