#!/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 getStrBuildtimeHMAC {
    my ( $self ) = @_;
    my $res_str = "";
    my $__buffer  = $self->{_buildtime_hmac};
    foreach(split(//, $__buffer)) {
        $res_str = sprintf("%s%02X:", $res_str, ord($_));
    }
    return $res_str;
}

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 searchPattern {
    my ( $string, $substring, $counter ) = ( shift, shift, \shift );
    my $offset;
    my $result;
    my $prev_result;

    if (!$$counter) {
        $$counter = 0;
    }

    $result = index($string, $substring);
    while ($result != -1) {
        $$counter++;
        $offset = $result + 1;
        $prev_result = $result;
        $result = index($string, $substring, $offset);
    }
    return $prev_result ? $prev_result : $result;
}

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_str;
    my $elf_size = -s $self->{_elf_file};

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

    read($self->{_FD}, $elf_str, $elf_size);

    # perform sections search
    $self->{_hmac_position}      = searchPattern($elf_str, $self->{_hmac_pattern},      $hmac_occur);
    $self->{_text_position}[0]   = searchPattern($elf_str, $self->{_text_pattern}[0],   @sect_occur[0]);
    $self->{_text_position}[1]   = searchPattern($elf_str, $self->{_text_pattern}[1],   @sect_occur[1]);
    $self->{_rodata_position}[0] = searchPattern($elf_str, $self->{_rodata_pattern}[0], @sect_occur[2]);
    $self->{_rodata_position}[1] = searchPattern($elf_str, $self->{_rodata_pattern}[1], @sect_occur[3]);

    # 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_SCRYPTO_text_start      : 0x%X\n", $ptr_text_start;
        printf "[IMPRINT] FIPS_SCRYPTO_text_end        : 0x%X\n", $ptr_text_end;
        printf "[IMPRINT] FIPS_SCRYPTO_rodata_start    : 0x%X\n", $ptr_rodata_start;
        printf "[IMPRINT] FIPS_SCRYPTO_rodata_end      : 0x%X\n", $ptr_rodata_end;
        printf "[IMPRINT] FIPS_embedded_hmac           : 0x%X\n", $sig;
        printf "[IMPRINT] FIPS_SCRYPTO_text size       : %d\n",   $ptr_text_end - $ptr_text_start;
        printf "[IMPRINT] FIPS_SCRYPTO_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);
    }

    if (defined $ENV{IN_OUT_HMAC_FILE} and $ENV{IN_OUT_HMAC_FILE}) {
        open(my $in_out_hmac_file, ">", $ENV{IN_OUT_HMAC_FILE}) or die "$!";
        print $in_out_hmac_file $imprint->getStrBuildtimeHMAC();
        close($in_out_hmac_file);
    }

    exit(0);
}

__END__