HtmlSteg/htmlsteg.py

#!/usr/bin/python
# -*- coding: utf-8 -*-#

# Copyright 2017 Weber Yann <yann.weber@member.fsf.org>
#
# This file is part of HtmlSteg.
#
#        HtmlSteg 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
#        any later version.
#
#        HtmlSteg 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 HtmlSteg.  If not, see <http://www.gnu.org/licenses/>.
#

import sys
import argparse

from steg import clean_html, hide, unhide

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Html steganography program')
    parser.add_argument('-u', '--unhide', action='store_const',
        default=False, const=True,
        help="Unhide from input")
    parser.add_argument('-i', '--input', metavar='FILE', type=str,
        default='-', help='Filename or - for stdin')
    parser.add_argument('-o', '--output', metavar='FILE', type=str,
        default='-', help='Filename or - for stdout')
    parser.add_argument('-s', '--hide', metavar='FILE', type=str,
        default='', help="Html filename to hide input in")
    args = parser.parse_args()

    if len(args.hide) > 0 and args.unhide:
        sys.stderr.write('--unhide and --hide cannot be set at same time')
        parser.print_help()
        exit(1)

    try:
        if args.output == '-':
            fdout = sys.stdout
        else:
            fdout = open(args.output, 'w+')
        if args.input == '-':
            data_in = sys.stdin.read()
        else:
            with open(args.input, 'r') as fdin:
                data_in = fdin.read()
    except Exception as e:
        parser.print_help()
        raise e

    if args.unhide:
        res = unhide(data_in)
        fdout.write(res)
    elif len(args.hide) > 0:
        try:
            with open(args.hide, 'r') as fdhtml:
                html = fdhtml.read()
        except Exception as e:
            parser.print_help()
            raise e

        data_in = bytearray(data_in)
        html = clean_html(html)
        html = hide(html, data_in)
        fdout.write(html)
    else:
        parser.print_help()
        exit(1)