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

use strict;
use FindBin;
use lib $FindBin::RealBin;
use POSIX qw/EXIT_FAILURE EXIT_SUCCESS strftime/;
use English qw/-no_match_vars/;
use IO::File;
use Crypt::OpenSSL::RSA;
use Crypt::OpenSSL::X509 qw/FORMAT_ASN1 FORMAT_PEM/;
use Convert::ASN1;
use Math::BigInt;
use File::Basename qw/basename/;
use Getopt::Long qw//;
use TLV::Tags qw/:header :digest :record :function/;
use TLV::Parser;
use TLV::Builder;

sub parse_tlv {
    my $tlv_file = shift;

    my ($file, $content);

    unless ($file = IO::File->new ($tlv_file, '<:raw')) {
        die 'Unable to read ' . $tlv_file . ': ' . $OS_ERROR;
    }

    $content = do {local $INPUT_RECORD_SEPARATOR; $file->getline};
    $file->close;

    my $parser = TLV::Parser->new ($content);

    # Header
    die 'Not a version tag: ' . $parser->tag if ($parser->next_tag != HEADER_VERSION);
    die 'Wrong version length: ' . $parser->length if ($parser->next_length != 2);

    my $version = join ('.', unpack ('CC', $parser->next_value));

    die 'Not a header_length tag: ' . $parser->tag if ($parser->next_tag != HEADER_LENGTH);
    die 'Wrong header_length length: ' . $parser->length if ($parser->next_length != 2);

    my $header_length = unpack ('S>', $parser->next_value);

    print 'Version:                  ', $version, "\n",
          'Header Length:            ', $header_length, ' bytes', "\n";

    my ($header_serial_number, $header_digest_algorithm, $header_signature_index, $header_signature_length);

    while ($parser->index < $header_length) {
        $parser->next_tag;
        next if ($parser->tag == HEADER_PADDING);
        $parser->next_length;

        if ($parser->tag == HEADER_SIGNER_ID) {
            my $signer_id = $parser->length;

            print 'Signer ID:                ', $signer_id, "\n";
        } elsif ($parser->tag == HEADER_SIGNER_NAME) {
            my $signer_name = unpack ('Z*', $parser->next_value);

            print 'Signer Name:              ', $signer_name, "\n";
        } elsif ($parser->tag == HEADER_SERIAL_NUMBER) {
            my $serial_number = uc unpack ('H*', $parser->next_value);

            # Used to locate the signing certificate record
            $header_serial_number = $serial_number;

            print 'Serial Number:            ', join (':', $serial_number =~ m/(..)/g), "\n";
        } elsif ($parser->tag == HEADER_CA_NAME) {
            my $ca_name = unpack ('Z*', $parser->next_value);

            print 'CA Name:                  ', $ca_name, "\n";
        } elsif ($parser->tag == HEADER_SIGNATURE_INFO) {
            my $signature_info = $parser->length;

            print 'Signature Info:           ', $signature_info, "\n";
        } elsif ($parser->tag == HEADER_DIGEST_ALGORITHM) {
            die 'Invalid digest_algorithm length: ' . $parser->length if ($parser->length != 1);

            my $digest_algorithm = unpack ('C', $parser->next_value);

            $header_digest_algorithm = $digest_algorithm;

            print 'Digest Algorithm:         ';

            if ($digest_algorithm == DIGEST_SHA1) {
                print 'SHA1';
            } elsif ($digest_algorithm == DIGEST_SHA256) {
                print 'SHA256';
            }

            print "\n";
        } elsif ($parser->tag == HEADER_SIGNATURE_ALGORITHM_INFO) {
            my $signature_algorithm_info = $parser->length;

            print 'Signature Algorithm Info: ', $signature_algorithm_info, "\n";
        } elsif ($parser->tag == HEADER_SIGNATURE_ALGORITHM) {
            die 'Invalid signature_algorithm length: ' . $parser->length if ($parser->length != 1);

            my $signature_algorithm = unpack ('C', $parser->next_value);

            print 'Signature Algorithm:      ', $signature_algorithm, "\n";
        } elsif ($parser->tag == HEADER_SIGNATURE_MODULUS) {
            die 'Invalid signature_modulus length: ' . $parser->length if ($parser->length != 1);

            my $signature_modulus = unpack ('C', $parser->next_value);

            print 'Signature Modulus:        ', $signature_modulus, "\n";
        } elsif ($parser->tag == HEADER_SIGNATURE) {
            my $signature = $parser->next_value;

            # The removal index for the signature 
            $header_signature_index  = $parser->index - $parser->length - 3;
            $header_signature_length = $parser->length + 3;

            print 'Signature:                ', length ($signature), ' bytes', "\n";
        } elsif ($parser->tag == HEADER_FILENAME) {
            my $filename = unpack ('Z*', $parser->next_value);

            print 'Filename:                 ', $filename, "\n";
        } elsif ($parser->tag == HEADER_TIMESTAMP) {
            die 'Invalid timestamp length: ' . $parser->length if ($parser->length != 4);

            my $timestamp = unpack ('L>', $parser->next_value);

            print 'Timestamp:                ', strftime ('%a, %d %b %Y %H:%M:%S %z', localtime $timestamp), "\n";
        } else {
            die 'Unknown tag: ' . $parser->tag . ' at index: ' . ($parser->index - 3);
        }
    }

    print "\n";

    die 'No header serial number' unless (length $header_serial_number);
    die 'No header digest algorithm' unless ($header_digest_algorithm);
    die 'No header signature' unless ($header_signature_index);

    # Records
    my ($record_number, $record_certificate);

    $record_number = 0;

    until ($parser->done) {
        die 'Not a record_length tag: ' . $parser->tag if ($parser->next_tag != RECORD_LENGTH);
        die 'Wrong record_length length: ' . $parser->length if ($parser->next_length != 2);

        my $record_length = unpack ('S>', $parser->next_value);
        my $record_index  = $parser->index - $parser->length - 3;

        $record_number++;

        print 'Record Number:            ', $record_number, "\n",
              'Record Length:            ', $record_length, ' bytes', "\n";

        my ($record_function, $record_serial_number);

        while ($parser->index - $record_index < $record_length) {
            $parser->next_tag;
            $parser->next_length;

            if ($parser->tag == RECORD_DNS_NAME) {
                my $dns_name = unpack ('Z*', $parser->next_value);

                print 'DNS Name:                 ', $dns_name, "\n";
            } elsif ($parser->tag == RECORD_SUBJECT_NAME) {
                my $subject_name = unpack ('Z*', $parser->next_value);

                print 'Subject Name:             ', $subject_name, "\n";
            } elsif ($parser->tag == RECORD_FUNCTION) {
                my $function = unpack ('S>', $parser->next_value);

                $record_function = $function;

                print 'Function:                 ';

                if ($function == FUNCTION_SAST) {
                    print 'SAST';
                } elsif ($function == FUNCTION_CCM) {
                    print 'CCM';
                } elsif ($function == FUNCTION_CCM_TFTP) {
                    print 'CCM+TFTP';
                } elsif ($function == FUNCTION_TFTP) {
                    print 'TFTP';
                } elsif ($function == FUNCTION_HTTPS) {
                    print 'HTTPS';
                }

                print "\n";
            } elsif ($parser->tag == RECORD_ISSUER_NAME) {
                my $issuer_name = unpack ('Z*', $parser->next_value);

                print 'Issuer Name:              ', $issuer_name, "\n";
            } elsif ($parser->tag == RECORD_SERIAL_NUMBER) {
                my $serial_number = uc unpack ('H*', $parser->next_value);

                $record_serial_number = $serial_number;

                print 'Serial Number:            ', join (':', $serial_number =~ m/(..)/g), "\n";
            } elsif ($parser->tag == RECORD_PUBLIC_KEY) {
                my $public_key = $parser->next_value;

                print 'Public Key:               ', length ($public_key), ' bytes', "\n";
            } elsif ($parser->tag == RECORD_SIGNATURE) {
                my $signature = $parser->next_value;

                print 'Signature:                ', length ($signature), ' bytes', "\n";
            } elsif ($parser->tag == RECORD_CERTIFICATE) {
                my $certificate = $parser->next_value;

                # Records that have the same serial number as the header and SAST function are used for signing
                if ($record_serial_number eq $header_serial_number && defined ($record_function) && $record_function == FUNCTION_SAST) {
                    $record_certificate = $certificate;
                }

                print 'Certificate:              ', length ($certificate), ' bytes', "\n";
            } else {
                die 'Invalid record tag: ' . $parser->tag . ' at index: ' . ($parser->index - 3);
            }
        }

        print "\n";
    }

    die 'No record certificate' unless (length $record_certificate);

    my $x509 = Crypt::OpenSSL::X509->new_from_string ($record_certificate, FORMAT_ASN1);
    die 'Unable to load record x509 certificate' unless ($x509);

    my $rsa = Crypt::OpenSSL::RSA->new_public_key ($x509->pubkey);
    die 'Unable to parse RSA public key' unless ($rsa);

    if ($header_digest_algorithm == DIGEST_SHA1) {
        $rsa->use_sha1_hash;
    } elsif ($header_digest_algorithm == DIGEST_SHA256) {
        $rsa->use_sha256_hash;
    } else {
        die 'Unknown header_digest_algorithm: ' . $header_digest_algorithm;
    }

    $content = $parser->content;
    my $signature = substr ($content, $header_signature_index + 3, $header_signature_length - 3);

    # Remove the signature block
    substr ($content, $header_signature_index, $header_signature_length, '');

    if ($rsa->verify ($content, $signature)) {
        print 'Valid signature', "\n";
    } else {
        print 'Invalid signature', "\n";
    }
}

sub build_tlv {
    my ($tlv_file, $certificate_file, $digest_algorithm, $filename, $records);

    $tlv_file         = shift;
    $certificate_file = shift;
    $digest_algorithm = shift;
    $filename         = shift;
    $records          = shift;

    # Automatically grant the signing certificate the SAST function
    unshift (@{$records}, {certificate_file => $certificate_file, function => 'SAST'});

    my ($file, $content);

    unless ($file = IO::File->new ($certificate_file, '<:raw')) {
        die 'Unable to read ' . $certificate_file . ': ' . $OS_ERROR;
    }

    $content = do {local $INPUT_RECORD_SEPARATOR; $file->getline};
    $file->close;

    my $x509 = Crypt::OpenSSL::X509->new_from_string ($content, FORMAT_PEM);
    die 'Unable to load header certificate' unless ($x509);

    my $rsa = Crypt::OpenSSL::RSA->new_private_key ($content);
    die 'Unable to load header private key' unless ($rsa);

    if ($digest_algorithm eq 'SHA1') {
        $rsa->use_sha1_hash;
    } elsif ($digest_algorithm eq 'SHA256') {
        $rsa->use_sha256_hash;
    } else {
        die 'Unknown digest_algorithm: ' . $digest_algorithm;
    }

    my $builder = TLV::Builder->new;
    my $header_signature_index;

    do {
        # Header
        $builder->next_tag (HEADER_VERSION);
        $builder->next_length (2);
        $builder->next_value (pack ('CC', 1, 2));

        $builder->next_tag (HEADER_LENGTH);
        $builder->next_length (2);
        $builder->next_value (pack ('S>', 0));

        (my $signer_name = $x509->subject) =~ s/, /;/g;
        (my $ca_name = $x509->issuer) =~ s/, /;/g;
        my $serial_number = pack ('H*', $x509->serial);

        # Combined TLV length for signer_name, serial_number and ca_name
        my $signer_id = 3 + length ($signer_name) + 1 + 3 + length ($serial_number) + 3 + length ($ca_name) + 1;

        $builder->next_tag (HEADER_SIGNER_ID);
        $builder->next_length ($signer_id);

        $builder->next_tag (HEADER_SIGNER_NAME);
        $builder->next_length (length ($signer_name) + 1);
        $builder->next_value (pack ('Z*', $signer_name));

        $builder->next_tag (HEADER_SERIAL_NUMBER);
        $builder->next_length (length $serial_number);
        $builder->next_value ($serial_number);

        $builder->next_tag (HEADER_CA_NAME);
        $builder->next_length (length ($ca_name) + 1);
        $builder->next_value (pack ('Z*', $ca_name));

        $builder->next_tag (HEADER_SIGNATURE_INFO);
        $builder->next_length (15); # Unknown

        $builder->next_tag (HEADER_DIGEST_ALGORITHM);
        $builder->next_length (1);
        $builder->next_value (pack ('C', do {
            if ($digest_algorithm eq 'SHA1') {
                DIGEST_SHA1;
            } elsif ($digest_algorithm eq 'SHA256') {
                DIGEST_SHA256;
            } else {
                die 'Unknown digest algorithm: ' . $digest_algorithm;
            }
        }));

        $builder->next_tag (HEADER_SIGNATURE_ALGORITHM_INFO);
        $builder->next_length (8); # Unknown

        $builder->next_tag (HEADER_SIGNATURE_ALGORITHM);
        $builder->next_length (1);
        $builder->next_value (pack ('C', 0)); # Unknown

        my $signature_modulus;

        if ($rsa->size == 64) {
            $signature_modulus = 0;
        } elsif ($rsa->size == 128) {
            $signature_modulus = 1;
        } elsif ($rsa->size == 256) {
            $signature_modulus = 2;
        } elsif ($rsa->size == 512) {
            $signature_modulus = 3;
        } else {
            die 'Unsupported RSA key size: ' . $rsa->size;
        }

        $builder->next_tag (HEADER_SIGNATURE_MODULUS);
        $builder->next_length (1);
        $builder->next_value (pack ('C', $signature_modulus));

        # The insertion index of the signature
        $header_signature_index = $builder->index;

        $builder->next_tag (HEADER_FILENAME);
        $builder->next_length (length ($filename) + 1);
        $builder->next_value (pack ('Z*', $filename));

        $builder->next_tag (HEADER_TIMESTAMP);
        $builder->next_length (4);
        $builder->next_value (pack ('L>', time));

        # Header must be padded to 32-bit boundary
        while (($builder->index + 3 + $rsa->size) % 4) {
            $builder->next_tag (HEADER_PADDING);
        }

        # Signed content includes the length of the signature block
        $builder->length (8, $builder->index + 3 + $rsa->size);
    };

    # Records
    foreach my $record (@{$records}) {
        my ($certificate_file, $function) = ($record->{certificate_file}, $record->{function});

        my ($file, $content);

        unless ($file = IO::File->new ($certificate_file, '<:raw')) {
            die 'Unable to read ' . $certificate_file . ': ' . $OS_ERROR;
        }

        $content = do {local $INPUT_RECORD_SEPARATOR; $file->getline};
        $file->close;

        my $x509 = Crypt::OpenSSL::X509->new_from_string ($content, FORMAT_PEM);
        die 'Unable to load record certificate' unless ($x509);

        my $record_index = $builder->index;

        $builder->next_tag (RECORD_LENGTH);
        $builder->next_length (2);
        $builder->next_value (pack ('S>', 0));

        unless ($function eq 'SAST') {
            (my $dns_name) = ($x509->subject =~ m/\bCN=([^,]+)/);

            # DNS name is included even if empty
            $builder->next_tag (RECORD_DNS_NAME);
            $builder->next_length (length ($dns_name) + 1);
            $builder->next_value (pack ('Z*', $dns_name));
        }

        (my $subject_name = $x509->subject) =~ s/, /;/g;

        $builder->next_tag (RECORD_SUBJECT_NAME);
        $builder->next_length (length ($subject_name) + 1);
        $builder->next_value (pack ('Z*', $subject_name));

        $builder->next_tag (RECORD_FUNCTION);
        $builder->next_length (2);
        $builder->next_value (pack ('S>', do {
            if ($function eq 'SAST') {
                FUNCTION_SAST;
            } elsif ($function eq 'CCM') {
                FUNCTION_CCM;
            } elsif ($function eq 'CCM+TFTP') {
                FUNCTION_CCM_TFTP;
            } elsif ($function eq 'TFTP') {
                FUNCTION_TFTP;
            } elsif ($function eq 'HTTPS') {
                FUNCTION_HTTPS;
            } else {
                die 'Unknown record function: ' . $function;
            };
        }));

        (my $issuer_name = $x509->issuer) =~ s/, /;/g;

        $builder->next_tag (RECORD_ISSUER_NAME);
        $builder->next_length (length ($issuer_name) + 1);
        $builder->next_value (pack ('Z*', $issuer_name));

        my $serial_number = pack ('H*', $x509->serial);

        $builder->next_tag (RECORD_SERIAL_NUMBER);
        $builder->next_length (length $serial_number);
        $builder->next_value ($serial_number);

        my $asn1 = Convert::ASN1->new;

        $asn1->prepare ('SEQUENCE { modulus INTEGER, exponent INTEGER }');

        my $public_key = $asn1->encode ({modulus  => Math::BigInt->from_hex ($x509->modulus),
                                         exponent => hex $x509->exponent});

        $builder->next_tag (RECORD_PUBLIC_KEY);
        $builder->next_length (length $public_key);
        $builder->next_value ($public_key);

        my $signature = pack ('H*', $x509->sig_print);

        $builder->next_tag (RECORD_SIGNATURE);
        $builder->next_length (length $signature);
        $builder->next_value ($signature);

        my $certificate = $x509->as_string (FORMAT_ASN1);

        $builder->next_tag (RECORD_CERTIFICATE);
        $builder->next_length (length $certificate);
        $builder->next_value ($certificate);

        my $record_length = $builder->index - $record_index;

        $builder->length ($record_index + 3, $record_length);
    };

    $content = $builder->content;
    die 'Unable to sign content' unless (my $signature = $rsa->sign ($content));

    # Insert signature into content
    substr ($content, $header_signature_index, 0, pack ('CS>a*', HEADER_SIGNATURE, $rsa->size, $signature));

    unless ($file = IO::File->new ($tlv_file, '>:raw')) {
        die 'Unable to write ' . $tlv_file . ': ' . $OS_ERROR;
    }

    $file->print ($content);
    $file->close;

    print 'Built ', $tlv_file, "\n";
}

eval {
    my ($mode, $tlv_file, $certificate_file, $digest_algorithm, $filename, $records, $record, $show_help);

    $records = [];

    my $getopt = Getopt::Long::Parser->new;

    $getopt->configure (qw/no_ignore_case/);

    unless ($getopt->getoptions ('p|parse'         => sub {shift; $mode = 'parse'},
                                 'b|build'         => sub {shift; $mode = 'build'},
                                 'c|certificate=s' => sub {shift; $certificate_file = shift},
                                 'd|digest=s'      => sub {shift; $digest_algorithm = uc shift},
                                 'r|record=s'      => sub {shift; push (@{$records}, $record = {certificate_file => shift})},
                                 'f|function=s'    => sub {shift; $record->{function} = uc shift if $record},
                                 'F|filename=s'    => sub {shift; $filename  = shift},
                                 'h|help'          => sub {shift; $show_help = shift})) {
        die 'Error parsing options';
    }

    if ($show_help) {
        print 'Usage: ', basename ($PROGRAM_NAME), ' <FILE> [OPTIONS]', "\n",
              'Parse or build a .tlv file', "\n",
              "\n",
              '  -p --parse                 parse a .tlv file', "\n",
              '  -b --build                 build a .tlv file', "\n",
              '  -c --certificate <file>    certificate to use for verifying or signing', "\n",
              '  -d --digest <name>         signature digest (sha1, sha256)', "\n",
              '  -F --filename <name>       header filename in built .tlv file (optional)', "\n",
              '  -r --record <file>         additional record certificate', "\n",
              '  -f --function <name>       record function (sast, ccm, ccm+tftp tftp, https)', "\n",
              '  -h --help                  print this help and exit', "\n",
              "\n";

        return;
    }

    die 'No .tlv file specified' unless (length ($tlv_file = shift));

    if ($mode eq 'parse') {
        parse_tlv ($tlv_file);
    } elsif ($mode eq 'build') {
        die 'No signing certificate file specified' unless (length $certificate_file);

        $digest_algorithm = 'SHA1' unless (length $digest_algorithm);
        $filename = basename ($tlv_file) unless (length $filename);

        build_tlv ($tlv_file, $certificate_file, $digest_algorithm, $filename, $records);
    } else {
        die 'No mode specified, choose either --build, --parse or --help for available options';
    }
};

if (length $EVAL_ERROR) {
    $EVAL_ERROR =~ s/ at \S+ line \d+\.//;
    warn $EVAL_ERROR;

    exit EXIT_FAILURE;
}

exit EXIT_SUCCESS;