#!/usr/bin/env python3
#
# Copyright (c) 2020 Gareth Palmer <gareth.palmer3@gmail.com>
# This program is free software, distributed under the terms of
# the GNU General Public License Version 2.

import sys
import os
import os.path
import struct
import getopt
import traceback


HEADER_VERSION = 1
HEADER_LENGTH = 2


class ProgramError(Exception):
    pass


def strip_sgn_file(sgn_file):
    try:
        with open(sgn_file, 'rb') as file:
            tlv_data = file.read()

    except (PermissionError, FileNotFoundError, IsADirectoryError) as error:
        raise ProgramError(f'{error.strerror}: {error.filename}')

    tlv_index = 0

    (tlv_tag, tlv_length) = struct.unpack_from('> B H', tlv_data, tlv_index)
    tlv_index += 3

    if tlv_tag != HEADER_VERSION:
        raise ProgramError(f'Tag is not header version: {tlv_tag}')

    # Skip version
    tlv_index += tlv_length

    (tlv_tag, tlv_length) = struct.unpack_from('> B H', tlv_data, tlv_index)
    tlv_index += 3

    if tlv_tag != HEADER_LENGTH:
        raise ProgramError(f'Tag is not header length: {tlv_tag}')

    (header_length,) = struct.unpack_from('> H', tlv_data, tlv_index)
    tlv_index += tlv_length

    # Remove .sgn or .sha512
    output_file = sgn_file[:sgn_file.rindex('.')]

    try:
        with open(output_file, 'wb') as file:
            file.write(tlv_data[header_length:])

    except (PermissionError, FileNotFoundError, IsADirectoryError) as error:
        raise ProgramError(f'{error.strerror}: {error.filename}')

    print(f'Wrote {output_file}')


def main():
    try:
        short_options = 'H'
        long_options = ['help']

        sgn_file = None
        help = False

        try:
            options, arguments = getopt.gnu_getopt(sys.argv[1:], short_options, long_options)
        except getopt.GetoptError as error:
            raise ProgramError(error.msg[0].upper() + error.msg[1:] +
                               '. Try \'' + os.path.basename(sys.argv[0]) + ' --help\' for more information')

        for option, _ in options:
            if option in ('-H', '--help'):
                help = True

        if help:
            print('Usage: ' + os.path.basename(sys.argv[0]) + ' [OPTIONS] SGN-FILE\n'
                  'Remove .sgn header from a file.\n'
                  '\n'
                  '  -H, --help    print this help and exit\n')

            return

        if len(arguments):
            sgn_file = arguments[0]

        if sgn_file is None:
            raise ProgramError('No .sgn file specified')

        if not sgn_file.endswith(('.sgn', '.sha512')):
            raise ProgramError(f'File {sgn_file} does not end with .sgn or .sha512: {sgn_file}')

        strip_sgn_file(sgn_file)

    except ProgramError as error:
        print(str(error), file = sys.stderr)
        exit(1)

    except Exception:
        traceback.print_exc(file = sys.stderr)
        exit(1)

    exit(0)


if __name__ == '__main__':
    main()