#!/usr/bin/env perl

use strict;

use Encode;
use Env;

{
    package Utils;
    sub printBuffer {
        my $__preline = shift;
        my $__buffer  = shift;

        $__preline = join "", $__preline;
        printf "%s ", $__preline;

        foreach(split(//, $__buffer)) {
            printf("%02X:", ord($_));
        }
        print "\n";
    }

    sub hexToBin {
        my $str = shift;
        return pack("H*", $str);
    }
}

package Imprint;

use Digest::SHA qw(hmac_sha1);
use FileHandle;

# Imprint Constructor
# arg: path to elf file
sub new
{
    my $class = shift;
    my $self = {
        _elf_file => shift,
        _DEBUG    => shift,

        _FD => FileHandle->new,

        _text_pattern   => undef,
        _rodata_pattern => undef,
        _hmac_pattern   => undef,

        _text_position   => undef,
        _rodata_position => undef,
        _hmac_position   => undef,

        _buildtime_hmac => undef,
        _HMAC_LEN => 20,
        _compatibility_version => 2.1,

        _covered_area => 0,
    };

    # Construct
    my $legacy_mode;
    sysopen($self->{_FD}, $self->{_elf_file}, $legacy_mode?0:2) or die "$!";
    binmode($self->{_FD});

    print "[IMPRINT] Imprint compatibility : SCrypto $self->{_compatibility_version} version or higher \n";
    if ($self->{_DEBUG}) {
        print "[IMPRINT] File : $self->{_elf_file} \n";
    }

    bless $self, $class;
    return $self;
}

sub getCoveredArea {
    my ( $self ) = @_;
    return $self->{_covered_area};
}

sub setTextSectionPatterns {
    my ( $self, $start, $end ) = @_;
    $self->{_text_pattern}[0] = Utils::hexToBin($start) if defined($start) or die "$!";
    $self->{_text_pattern}[1] = Utils::hexToBin($end)   if defined($end)   or die "$!";
}

sub setRodataSectionPatterns {
    my ( $self, $start, $end ) = @_;
    $self->{_rodata_pattern}[0] = Utils::hexToBin($start) if defined($start) or die "$!";
    $self->{_rodata_pattern}[1] = Utils::hexToBin($end)   if defined($end)   or die "$!";
}

sub setHmacPattern {
    my ( $self, $pattern ) = @_;
    $self->{_hmac_pattern} = Utils::hexToBin($pattern) if defined($pattern) or die "$!";
}

sub isReady {
    my ( $self ) = @_;
    return defined($self->{_text_pattern}) && defined($self->{_rodata_pattern}) && defined($self->{_hmac_pattern})
}

sub findSections {
    my ( $self ) = @_;

    if (!isReady($self)) {
        print "[IMPRINT] Imprint object is not ready, please specify section boundarie patterns \n";
        return 0;
    }

    if ($self->{_DEBUG}) {
        print "[IMPRINT] Procedure started ... \n";
    }

    my $elf_size = -s $self->{_elf_file};
    my $sect_pattern_len = length ($self->{_text_pattern}[0]);
    my $hmac_pattern_len = length ($self->{_hmac_pattern});
    my $temp_read;
    my $hmac_read;

    my $hmac_occur   =  0;
    my @sect_occur   = (0) x 4; # text start, text end, rodata start, rodata end

    # Body of ELF file (search window size - sect_pattern_len)
    my $_i = 0;
    while ($_i < ($elf_size - $sect_pattern_len)) {
        seek($self->{_FD}, $_i, 0);
        read($self->{_FD}, $temp_read, $sect_pattern_len);

        $hmac_read = substr($temp_read, 0, $hmac_pattern_len);
        if ($self->{_hmac_pattern}           eq $hmac_read) { # hmac
            $self->{_hmac_position} = $_i;
            $hmac_occur++;
        } elsif ($self->{_text_pattern}[0]   eq $temp_read) { # start
            $self->{_text_position}[0] = $_i;
            @sect_occur[0]++;
        } elsif ($self->{_text_pattern}[1]   eq $temp_read) { # end
            $self->{_text_position}[1] = $_i;
            @sect_occur[1]++;
        } elsif ($self->{_rodata_pattern}[0] eq $temp_read) { # start
            $self->{_rodata_position}[0] = $_i;
            @sect_occur[2]++;
        } elsif ($self->{_rodata_pattern}[1] eq $temp_read) { # end
            $self->{_rodata_position}[1] = $_i;
            @sect_occur[3]++;
        }

        $_i++;
    }

    # Tail of ELF file (search window size - hmac_pattern_len)
    my $_j = 0;
    while ($_j < ($sect_pattern_len - $hmac_pattern_len)) {

        $hmac_read = substr($temp_read, $_j, $hmac_pattern_len);
        if ($self->{_hmac_pattern} eq $hmac_read) { # hmac in tail
            $self->{_hmac_position} = $_i + $_j;
            $hmac_occur++;
        }

        $_j++;
    }

    # evaluates the expression for each element of a list
    if (grep { $_ != 1 } @sect_occur) {
        print "[IMPRINT] Bad file. Inconsistent ELF file section pattern occurances \n";
        print "[IMPRINT] Make sure, you use SCrypto version $self->{_compatibility_version} or higher \n";
        exit(1);
    }

    if ($hmac_occur == 0) {
        print "[IMPRINT] Bad file. HMAC pattern is missing \n";
        print "[IMPRINT] Make sure, you use SCrypto version $self->{_compatibility_version} or higher \n";
        exit(1);
    }

    return 1;
}

sub calculateFipsHmac {
    my ( $self, $hmac_key ) = @_;

    if (!defined($hmac_key)) {
        print "[IMPRINT] Specify key for HMAC \n";
        exit(1);
    }

    if (!defined($self->{_text_position}) || !defined($self->{_rodata_position}) || !defined($self->{_hmac_position})) {
        if (!findSections($self)) {
            exit(1);
        }
    }

    if ($self->{_DEBUG}) {
        print "[IMPRINT] Calculating HMAC ... \n";
    }

    my $ptr_text_start   = $self->{_text_position}[0];
    my $ptr_text_end     = $self->{_text_position}[1];
    my $ptr_rodata_start = $self->{_rodata_position}[0];
    my $ptr_rodata_end   = $self->{_rodata_position}[1];
    my $sig              = $self->{_hmac_position};

    my $blob;
    my $temp_blob;

    if ($ptr_text_start <= $ptr_rodata_start && $ptr_text_end >= $ptr_rodata_start) {
        $ptr_rodata_start = $ptr_text_start;
        $ptr_rodata_end   = $ptr_text_end > $ptr_rodata_end ? $ptr_text_end : $ptr_rodata_end;
        $ptr_text_start   = $ptr_text_end = 0;
    } elsif ($ptr_rodata_start <= $ptr_text_start && $ptr_rodata_end >= $ptr_text_start) {
        $ptr_rodata_start = $ptr_rodata_start;
        $ptr_rodata_end   = $ptr_text_end > $ptr_rodata_end ? $ptr_text_end : $ptr_rodata_end;
        $ptr_text_start   = $ptr_text_end = 0;
    }

    if ($self->{_DEBUG}) {
        printf "[IMPRINT] FIPS_bssl_text_start      : 0x%X\n", $ptr_text_start;
        printf "[IMPRINT] FIPS_bssl_text_end        : 0x%X\n", $ptr_text_end;
        printf "[IMPRINT] FIPS_bssl_rodata_start    : 0x%X\n", $ptr_rodata_start;
        printf "[IMPRINT] FIPS_bssl_rodata_end      : 0x%X\n", $ptr_rodata_end;
        printf "[IMPRINT] FIPS_embedded_hmac        : 0x%X\n", $sig;
        printf "[IMPRINT] FIPS_bssl_text size       : %d\n",   $ptr_text_end - $ptr_text_start;
        printf "[IMPRINT] FIPS_bssl_rodata size     : %d\n",   $ptr_rodata_end - $ptr_rodata_start;
    }

    if ($ptr_text_start) {
        seek($self->{_FD}, $ptr_text_start, 0)                          or die "$!";
        read($self->{_FD}, $temp_blob, $ptr_text_end - $ptr_text_start) or die "$!";
        $blob = $blob . $temp_blob;

        if ($self->{_DEBUG}) {
            $self->{_covered_area} += ($ptr_text_end - $ptr_text_start);
        }
    }

    if ($sig >= $ptr_rodata_start && $sig < $ptr_rodata_end) {
        seek($self->{_FD}, $ptr_rodata_start, 0)                 or die "$!";
        read($self->{_FD}, $temp_blob, $sig - $ptr_rodata_start) or die "$!";
        $blob = $blob . $temp_blob;

        if ($self->{_DEBUG}) {
            $self->{_covered_area} += ($sig - $ptr_rodata_start);
        }

        $ptr_rodata_start = $sig + $self->{_HMAC_LEN};

        seek($self->{_FD}, $ptr_rodata_start, 0)                            or die "$!";
        read($self->{_FD}, $temp_blob, $ptr_rodata_end - $ptr_rodata_start) or die "$!";
        $blob = $blob . $temp_blob;

        if ($self->{_DEBUG}) {
            $self->{_covered_area} += ($ptr_rodata_end - $ptr_rodata_start);
        }
    } else {
        seek($self->{_FD}, $ptr_rodata_start, 0)                            or die "$!";
        read($self->{_FD}, $temp_blob, $ptr_rodata_end - $ptr_rodata_start) or die "$!";
        $blob = $blob . $temp_blob;

        if ($self->{_DEBUG}) {
            $self->{_covered_area} += ($ptr_rodata_end - $ptr_rodata_start);
        }
    }

    if ($self->{_DEBUG}) {
        print "[IMPRINT] Integrity Protection : ".$self->{_covered_area}." bytes\n";
    }

    $self->{_buildtime_hmac} = hmac_sha1($blob, $hmac_key);

    if ($self->{_DEBUG}) {
        Utils::printBuffer("[IMPRINT] HMAC generated : ", $self->{_buildtime_hmac});
    }
}

sub embedHmac {
    my( $self ) = @_;

    # Embed HMAC right after pattern
    seek(   $self->{_FD},  $self->{_hmac_position} + $self->{_HMAC_LEN}, 0) or die "$!";
    print { $self->{_FD} } $self->{_buildtime_hmac}                         or die "$!";

    if ($self->{_DEBUG}) {
        print "[IMPRINT] HMAC was embeded \n";
    }
}

sub DESTROY
{
    my ( $self ) = @_;
    close($self->{_FD});
}

########################################################################################################################

__PACKAGE__->run( @ARGV ) unless caller();

sub run {
    my( $elf_file ) = @_;

    if(!defined($elf_file)) {
        print STDERR "usage: $0 [ elf name ] binary\n";
        exit(1);
    }

    my $imprint  = new Imprint($ARGV[0], 1); # DEBUG
    $imprint->setTextSectionPatterns(  "FFFFFFFF0D65798D143EFF3340E7952B5D6BC6F9B4797647FFB35A7EA3E925F7",
                                       "FFFFFFFF4AD75EEF1B18E0C68DE717E1A8A52DF0120F321725A7F6E450B911D0");
    $imprint->setRodataSectionPatterns("752D3C58FD2E50346D2802F6FDCA3CE675DE2F06717A66368D5780B484BB781E",
                                       "40BB5D62BB0C076168E2F7CD06031699F6A24734E5B520FA9F7900748A8A6908");
    $imprint->setHmacPattern(          "DA6885BCAC2908A9B43499E73644D4187A2773BD");

    $imprint->calculateFipsHmac("W1thout nigh2mare, there'1 be no 6ream.");

    $imprint->embedHmac();

    if (defined $ENV{IN_OUT_FILE}) {
        open(my $in_out_file, ">", $ENV{"IN_OUT_FILE"}) or die "$!";
        print $in_out_file $imprint->getCoveredArea();
        close($in_out_file);
    }

    exit(0);
}

__END__