#!/usr/bin/env python3

"""
A command-line interface for the steganography program
Copyright (C) 2022  Valentin Moguérou

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""

from argparse import ArgumentParser as ArgParser
from pathlib import Path

from libstegano import read_file, write_file, decode_string, encode_string, hex_print

def read(args):
    if args.verbose:
        print("Read mode enabled.")
    
    with open(args.from_file, 'rb') as from_file:
        byte_list = read_file(from_file, verbose=args.verbose)

        if args.verbose:
            print(f"Read {(len(byte_list)+1)*8} bits, {len(byte_list)+1} bytes.")
            print("======== BEGIN MESSAGE ========")
            print(decode_string(byte_list))
            print("========= END MESSAGE =========")
        else:
            print(decode_string(byte_list))

def write(args):
    byte_list = encode_string(args.string)

    with open(args.from_file, 'rb') as from_file, open(args.to_file, 'wb') as to_file:
        if args.verbose:
            print("Write mode enabled.")
            print("================== BYTE STREAM ==================")
            hex_print(byte_list, margin=1)
            print("================== BYTE STREAM ==================")

        write_file(from_file, to_file, byte_list, verbose=args.verbose)

def extract(args):
    if args.verbose:
        print("Extract mode enabled.")
    
    with open(args.from_file, 'rb') as from_file, open(args.to_file, 'wb') as to_file:
        byte_list = read_file(from_file, verbose=args.verbose)

        if args.verbose:
            print(f"Read {(len(byte_list)+1)*8} bits, {len(byte_list)+1} bytes.")
            print("================== BYTE STREAM ==================")
            hex_print(byte_list, margin=1)
            print("================== BYTE STREAM ==================")
        
        to_file.write(byte_list)

def insert(args):
    with open(args.from_file, 'rb') as from_file, open(args.content_file, 'rb') as content_file, open(args.to_file, 'wb') as to_file:
        byte_list = content_file.read()+b'\0'

        if args.verbose:
            print("Insert mode enabled.")
            print("================== BYTE STREAM ==================")
            hex_print(byte_list, margin=1)
            print("================== BYTE STREAM ==================")

        write_file(from_file, to_file, byte_list, verbose=args.verbose)

def main():
    parser = ArgParser()
    parser.add_argument('-v', '--verbose', action='store_true')

    subparsers = parser.add_subparsers(required=True)

    parser_read = subparsers.add_parser('read', help='Retrieve data from an image')
    parser_read.add_argument('from_file', type=Path)
    parser_read.set_defaults(func=read)

    parser_write = subparsers.add_parser('write', help='Dissimulate a string in an image')
    parser_write.add_argument('string', type=str)
    parser_write.add_argument('to_file', type=Path)
    parser_write.add_argument('--from', type=Path, dest='from_file', default=Path('default_image.jpg'))
    parser_write.set_defaults(func=write)

    parser_read = subparsers.add_parser('extract', help='Retrieve data from an image')
    parser_read.add_argument('from_file', type=Path)
    parser_read.add_argument('to_file', type=Path)
    parser_read.set_defaults(func=extract)

    parser_write = subparsers.add_parser('insert', help='Dissimulate data in an image')
    parser_write.add_argument('content_file', type=Path)
    parser_write.add_argument('to_file', type=Path)
    parser_write.add_argument('--from', type=Path, dest='from_file', default=Path('default_image.jpg'))
    parser_write.set_defaults(func=insert)

    args = parser.parse_args()
    args.func(args)

if __name__ == "__main__":
    main()