sgnfile

#!/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 binascii
from datetime import datetime
import time
import getopt
import traceback

from cryptography import x509
from cryptography.exceptions import InvalidSignature
from cryptography.hazmat import backends
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, rsa


HEADER_VERSION = 1
HEADER_LENGTH = 2
HEADER_SIGNER_INFO = 3
HEADER_SIGNER_NAME = 4
HEADER_SERIAL_NUMBER = 5
HEADER_ISSUER_NAME = 6
HEADER_SIGNATURE_INFO = 7
HEADER_HASH_ALGORITHM = 8
HEADER_SIGNATURE_ALGORITHM_INFO = 9
HEADER_SIGNATURE_ALGORITHM = 10
HEADER_SIGNATURE_MODULUS = 11
HEADER_SIGNATURE = 12
HEADER_PADDING = 13
HEADER_FILENAME = 14
HEADER_TIMESTAMP = 15

HASH_SHA1 = 1
HASH_SHA256 = 2
HASH_SHA512 = 3


class ProgramError(Exception):
    pass


def parse_sgn_file(sgn_file, tftp_certificate_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}')

    (major, minor) = struct.unpack_from('B B', tlv_data, tlv_index)
    tlv_index += tlv_length

    version = f'{major}.{minor}'
    print(f'Version:          {version}')

    (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

    print(f'Header Length:    {header_length} bytes')

    hash_algorithm = None

    signature_index = 0
    signature_length = 0

    while tlv_index < header_length:
        tlv_tag = tlv_data[tlv_index]
        tlv_index += 1

        if tlv_tag == HEADER_PADDING:
            continue

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

        if tlv_tag in (HEADER_SIGNER_INFO, HEADER_SIGNATURE_INFO, HEADER_SIGNATURE_ALGORITHM_INFO):
            continue

        if tlv_tag == HEADER_SIGNER_NAME:
            signer_name = tlv_data[tlv_index:tlv_index + tlv_length]
            signer_name = signer_name[:-1].decode('utf-8')

            print(f'Signer Name:      {signer_name}')

        elif tlv_tag == HEADER_SERIAL_NUMBER:
            serial_number = binascii.hexlify(tlv_data[tlv_index:tlv_index + tlv_length])
            serial_number = serial_number.decode('utf-8')

            print(f'Serial Number:    {serial_number}')

        elif tlv_tag == HEADER_ISSUER_NAME:
            issuer_name = tlv_data[tlv_index:tlv_index + tlv_length]
            issuer_name = issuer_name[:-1].decode('utf-8')

            print(f'Issuer Name:      {issuer_name}')

        elif tlv_tag == HEADER_HASH_ALGORITHM:
            hash_algorithm = tlv_data[tlv_index]

            print('Digest Algorithm: ', end = '')

            if hash_algorithm == HASH_SHA1:
                print('SHA1')
            elif hash_algorithm == HASH_SHA512:
                print('SHA512')
            else:
                print(f'{hash_algorithm}')

        elif tlv_tag in (HEADER_SIGNATURE_ALGORITHM, HEADER_SIGNATURE_MODULUS):
            pass

        elif tlv_tag == HEADER_SIGNATURE:
            signature = tlv_data[tlv_index:tlv_index + tlv_length]

            print(f'Signature:        {tlv_length} bytes')

            signature_index = tlv_index - 3
            signature_length = tlv_length

        elif tlv_tag == HEADER_FILENAME:
            filename = tlv_data[tlv_index:tlv_index + tlv_length]
            filename = filename[:-1].decode('utf-8')

            print(f'Filename:         {filename}')

        elif tlv_tag == HEADER_TIMESTAMP:
            (timestamp,) = struct.unpack_from('> I', tlv_data, tlv_index)
            timestamp = datetime.fromtimestamp(timestamp)

            print(timestamp.strftime('Timestamp:        %Y-%m-%d %H:%M:%S'))

        else:
            raise ProgramError(f'Unknown HEADER tag: {tlv_tag}')

        tlv_index += tlv_length

    print('')

    if hash_algorithm is None:
        raise ProgramError('Missing HEADER_HASH_ALGORITHM')

    if hash_algorithm not in (HASH_SHA1, HASH_SHA512):
        raise ProgramError(f'Unsupported HEADER_HASH_ALGORITHM: {hash_algorithm}')

    if signature_index is None:
        raise ProgramError('Missing HEADER_SIGNATURE')

    if tftp_certificate_file is None:
        print('No certificate specified, unable to check signature')
        return

    try:
        with open(tftp_certificate_file, 'rb') as file:
            certificate = file.read()

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

    try:
        certificate = x509.load_pem_x509_certificate(certificate, backends.default_backend())
    except ValueError:
        raise ProgramError(f'No certificate in file: {tftp_certificate_file}')

    public_key = certificate.public_key()

    signature = tlv_data[signature_index + 3:signature_index + 3 + signature_length]
    tlv_data = tlv_data[:signature_index] + tlv_data[signature_index + 3 + signature_length:]

    try:
        public_key.verify(signature, tlv_data, padding.PKCS1v15(), hashes.SHA512() if hash_algorithm == HASH_SHA512 else hashes.SHA1())
    except InvalidSignature:
        print('Invalid signature')
    else:
        print('Valid signature')


def remove_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 build_sgn_file(input_file, tftp_certificate_file, hash_algorithm, filename):
    try:
        with open(tftp_certificate_file, 'rb') as file:
            certificate = private_key = file.read()

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

    try:
        certificate = x509.load_pem_x509_certificate(certificate, backends.default_backend())
    except ValueError:
        raise ProgramError(f'No certificate in file: {tftp_certificate_file}')

    public_key = certificate.public_key()

    if not isinstance(public_key, rsa.RSAPublicKey):
        raise ProgramError(f'No RSA public-key in file: {tftp_certificate_file}')

    try:
        private_key = serialization.load_pem_private_key(private_key, None, backends.default_backend())
    except ValueError:
        raise ProgramError(f'No private-key in file: {tftp_certificate_file}')

    if not isinstance(private_key, rsa.RSAPrivateKey):
        raise ProgramError(f'No RSA private-key in file: {tftp_certificate_file}')

    signature_length = private_key.key_size // 8

    try:
        with open(input_file, 'rb') as file:
            input = file.read()

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

    sgn_file = input_file + '.sgn'

    tlv_data = bytearray()
    tlv_data += struct.pack('> B H B B', HEADER_VERSION, 2, 1, 0)

    header_length_index = len(tlv_data)
    tlv_data += struct.pack('> B H H', HEADER_LENGTH, 2, 0)

    signer_name = ''

    for attribute in certificate.subject:
        signer_name += (';' if len(signer_name) else '') + attribute.rfc4514_string()

    signer_name = signer_name.encode('utf-8') + b'\x00'
    issuer_name = ''

    for attribute in certificate.issuer:
        issuer_name += (';' if len(issuer_name) else '') + attribute.rfc4514_string()

    issuer_name = issuer_name.encode('utf-8') + b'\x00'

    serial_number = certificate.serial_number
    serial_number = serial_number.to_bytes((serial_number.bit_length() + 7) // 8, byteorder = 'big')

    signer_info = 3 + len(signer_name) + 3 + len(serial_number) + 3 + len(issuer_name)
    tlv_data += struct.pack('> B H', HEADER_SIGNER_INFO, signer_info)

    tlv_data += struct.pack('> B H', HEADER_SIGNER_NAME, len(signer_name))
    tlv_data += signer_name

    tlv_data += struct.pack('> B H', HEADER_SERIAL_NUMBER, len(serial_number))
    tlv_data += serial_number

    tlv_data += struct.pack('> B H', HEADER_ISSUER_NAME, len(issuer_name))
    tlv_data += issuer_name

    tlv_data += struct.pack('> B H', HEADER_SIGNATURE_INFO, 15)
    tlv_data += struct.pack('> B H B', HEADER_HASH_ALGORITHM, 1, HASH_SHA512 if hash_algorithm == 'sha512' else HASH_SHA1)

    tlv_data += struct.pack('> B H', HEADER_SIGNATURE_ALGORITHM_INFO, 8)
    tlv_data += struct.pack('> B H B', HEADER_SIGNATURE_ALGORITHM, 1, 0)
    tlv_data += struct.pack('> B H B', HEADER_SIGNATURE_MODULUS, 1, [64, 128, 256, 512].index(signature_length))

    signature_index = len(tlv_data)
    filename = filename.encode('utf-8') + b'\x00'

    tlv_data += struct.pack('> B H', HEADER_FILENAME, len(filename))
    tlv_data += filename

    tlv_data += struct.pack('> B H I', HEADER_TIMESTAMP, 4, int(time.time()))

    # Pad to 4 byte boundary
    while (len(tlv_data) + 3 + signature_length) % 4:
        tlv_data.append(HEADER_PADDING)

    header_length = len(tlv_data) + 3 + signature_length
    struct.pack_into('> H', tlv_data, header_length_index + 3, header_length)

    tlv_data += input

    tlv_data = bytes(tlv_data)
    signature = private_key.sign(tlv_data, padding.PKCS1v15(), hashes.SHA512() if hash_algorithm == 'sha512' else hashes.SHA1())

    try:
        with open(sgn_file, 'wb') as file:
            file.write(tlv_data[:signature_index])
            file.write(struct.pack('> B H', HEADER_SIGNATURE, len(signature)) + signature)
            file.write(tlv_data[signature_index:])

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

    print(f'Built {sgn_file}')


def main():
    try:
        short_options = 'pbrd:F:t:H'
        long_options = ['parse', 'build', 'remove', 'digest=', 'filename=', 'tftp=', 'help']

        mode = None
        sgn_file = input_file = None
        tftp_certificate_file = None
        hash_algorithm = 'sha1'
        filename = 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, argument in options:
            if option in ('-p', '--parse'):
                mode = 'parse'

            elif option in ('-b', '--build'):
                mode = 'build'

            elif option in ('-r', '--remove'):
                mode = 'remove'

            elif option in ('-d', '--digest'):
                hash_algorithm = argument

                if hash_algorithm not in ('sha1', 'sha512'):
                    raise ProgramError(f'Invalid digest: {hash_algorithm}')

            elif option in ('-F', '--filename'):
                filename = argument

            elif option in ('-t', '--tftp'):
                tftp_certificate_file = argument

            elif option in ('-H', '--help'):
                help = True

        if help:
            print('Usage: ' + os.path.basename(sys.argv[0]) + ' [OPTIONS] {SGN-FILE | FILE}\n'
                  'Parse, build or remove (unsign) .sgn files.\n'
                  '\n'
                  '  -p, --parse                    parse the specified .sgn file\n'
                  '  -b, --build                    build new .sgn file using options below\n'
                  '  -r, --remove                   remove (unsign) a .sgn file\n'
                  '  -d, --digest ALGORITHM         digest algorithm for signature: sha1 or sha512 (default sha1)\n'
                  '  -F, --filename NAME            file name to store in .sgn header\n'
                  '  -t, --tftp TFTP-CERT-FILE      signing/verifying certificate with TFTP role\n'
                  '  -H, --help                     print this help and exit\n')

            return

        if mode is None:
            raise ProgramError('No mode specified (either parse, build or remove). Try --help')

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

        if mode in ('parse', 'remove'):
            if sgn_file is None:
                raise ProgramError('No .sgn or .sha512 file specified')

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

            if mode == 'parse':
                parse_sgn_file(sgn_file, tftp_certificate_file)
            elif mode == 'remove':
                remove_sgn_file(sgn_file)

        elif mode == 'build':
            if input_file is None:
                raise ProgramError('No file specified')

            if tftp_certificate_file is None:
                raise ProgramError('No certificate file specified')

            if filename is None:
                filename = os.path.basename(input_file) + '.sgn'

            build_sgn_file(input_file, tftp_certificate_file, hash_algorithm, filename)

    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()