#perl -w
# (C) 2003-2007 Willem Jan Hengeveld <itsme@xs4all.nl>
# Web: http://www.xs4all.nl/~itsme/
#      http://wiki.xda-developers.com/
#
# $Id: $
#
use strict;
use warnings;
use bigint lib=>'GMP';
use Crypt::Rijndael;
use Digest::SHA qw(sha256);
use IO::File;
use Getopt::Long;
$|=1;

# perl ~/cvsprj/xda-devtools/xda2nbftool/startrek_cidedit.pl -i ffffffff21fffffff74f0fc30700fde2 ~/phones/htc_herald/devices/35244201643029801/A-cid.nb -t herald
# perl ~/cvsprj/xda-devtools/xda2nbftool/startrek_cidedit.pl -i ffffffff27ffffff7f571fd507007b9b ~/phones/htc_herald/devices/35244201694667301/A-cid.nb -t herald

# perl ~/cvsprj/xda-devtools/xda2nbftool/startrek_cidedit.pl -i ec003fed82f53334313531371e1e0408 ~/phones/htc_viva/devices/x-lock.nb -t viva

# perl ~/cvsprj/xda-devtools/xda2nbftool/startrek_cidedit.pl -i 000000000909011515250bac180a054e ~/phones/htc_startrek/devices/00171/bdk1-0-cid.bin
# perl ~/cvsprj/xda-devtools/xda2nbftool/startrek_cidedit.pl -i 000000008b0601110003092d0f0206cf ~/phones/htc_startrek/devices/248547/bdk1-0-cid.bin
# perl ~/cvsprj/xda-devtools/xda2nbftool/startrek_cidedit.pl -i 000000008805010a2102102b090c05cd ~/phones/htc_startrek/devices/71329/bdk1-0-cid.bin


#  exp*y-b*(p-1)*(q-1)=1
my %devinfo= (
startrek=>{
    size=> 0x10000,
# pq = 320353134631681309556473161605063339987 * 276065138433109647164739983515002720659
    pq  => 88438332459575713003174419268502556144700924300208596763704599584326505691433,
    y   => 69399065592429616374372553429979240083906900467887630575728849318636934515345,
    exp => 65537,
},
herald=>{
# pq = 277743565283690073521434740970457506457*290570584256061099044034087541094874209
    size=> 0x10000,
    pq => 80704110037843272917076340670641752126086650004317299342742802800530220267513,
    y  => 37685397248701003724328047569215092852850109589449582779985915613251862205185,
    exp=> 65537,
},
viva=>{
    size=> 0x20000,
# pq = 290123772068167912897153496578933816457*332396879999647872961510943295669884507
    pq => 96436236649188001138471945066857677842570427201821167402703570072037925931699,
    y  => 26839753978381511829435712315478035977190439224674628583719899884502784648193,
    exp=> 65537,
},
);

my $docidhex;
my $newcid;
my $newcidfile;
my $clearlocks;
my $newstr;
my $dev= $devinfo{startrek};
#perl tstrsa.pl -i 000000000909011515250bac180a054e ..\devices\00171\cid.nb

# layout of cid data block:
# 04b5 : index to lockdatakey
# 07ba : index to lockdata block
# 1000-9000 : contains 2048 byte aes encrypted lockdata block
# b000-f000 : contains 64 bytes rsa encrypted lockdata key
# ffe0-END  : aes encrypted hash of cid data block

# layout of the encrypted data block:
# 0000-0008 : ?
# 0008-000c : length of cid
# 000c-?    : cid string
# 002c-0034 : str
# 0404-040c : probably sim lock
# 07e0-0800 : sha256 of 0000-07e0
sub usage {
    return <<__EOF__
Usage: startrek_cidedit [-l] [-c NEWCID] [-s STR] [-o out.nb] [-t devname] -i DOCID  cidfile.nb
__EOF__
}
GetOptions(
    "i=s" => \$docidhex,
    "c=s" => \$newcid,
    "o=s" => \$newcidfile,
    "s=s" => \$newstr,
    "t=s" => sub { $dev= $devinfo{$_[1]}; },
    "l" => \$clearlocks,
) or die usage();

if (!@ARGV || !$docidhex) {
    die usage();
}

my $ifn= shift;
my $ifh= IO::File->new($ifn, "r") or die "$ifn: $!\n";
binmode $ifh;

my $ciddata;
$ifh->read($ciddata, $dev->{size});
$ifh->close();

#my $docid= pack("H*", "000000000909011515250bac180a054e");
#my $docid= pack("H*", "000000008805010a2102102b090c05cd");

my $docid= pack("H*", $docidhex);

print "decrypting\n";
my $lockdata= decrypt_cid_block($ciddata, $docid);

#print "oldcid\n";
#print hexdump($newcid);
#print "lockdata\n";
#print hexdump($lockdata);

my $changed= 0;
if ($newcid) {
    substr($lockdata, 12,8)=$newcid;
    substr($lockdata, 8, 2)= pack("v", length($newcid));
    $changed=1;
}
if ($clearlocks) {
    substr($lockdata, 0, 8)= pack("Vvv", 1, 0, 0);
    substr($lockdata, 10, 2)= pack("v", 0);
    substr($lockdata, 0x0404, 8)= pack("VV", 0, 0);
    $changed=1;
}
if ($newstr) {
    $newstr = substr($newstr, 0, 16) if length($newstr)>16;
    $newstr .= "0" x (16-length($newstr)) if length($newstr)<16;

    $newstr =~ s/(\d)(\d)/$2$1/g;
    substr($lockdata, 0x2c, 8)= pack("H*", $newstr);
    $changed=1;
}
if ($changed && !$newcidfile) {
    die "need output filename\n";
}
if ($changed) {
    print "re-encrypting\n";
    my $newciddata= encrypt_cid_block($ciddata, $lockdata, $docid);

    my $ofh= IO::File->new($newcidfile, "w") or die "$newcidfile: $!\n";
    binmode $ofh;

    $ofh->print($newciddata);
    $ofh->close();
}
else {
    print "lockdata\n";
    print hexdump($lockdata);
}
#print "re-decrypting\n";
#my $lockdata2= decrypt_cid_block($newcid, $docid);
#print "newcid\n";
#print hexdump($newcid);
#print "lockdata2\n";
#print hexdump($lockdata2);
exit(0);

sub decrypt_cid_block {
    my ($ciddata, $docid)= @_;
    # part1: verify cid block checksum

    my $shadoc = swapdw(sha256($docid));
    #printf("decrypt_cid_block\n");

    # these are used to calculate the aes keys
    my $sha64k = swapdw(sha256(substr($ciddata, 0, -0x20)));
    my $bufend = substr($ciddata, -0x20, 0x20);

    my $aeskey1= substr($shadoc, 0, 16) . substr($sha64k, 0, 16);
    my $aesiv1= substr($sha64k, 16, 16) . substr($shadoc, 0, 16);

    my $result= doaesdecrypt($bufend, $aeskey1, $aesiv1);

    if ($result ne $sha64k) {
        printf("bufend=%s\n", unpack("H*", $bufend));
        printf("result=%s\n", unpack("H*", $result));
        printf("sha64k=%s\n", unpack("H*", $sha64k));
        die "buffer checksum failed\n";
    }

    # part2: decrypt the data section

    my $ciddata_4b5= ord(substr($ciddata, 0x4b5, 1));

    my $offset_4b5= 0xB000 + (($ciddata_4b5 &0xf)<<10);
    #printf("ciddata[4b5]= %02x -> offset %04x\n", $ciddata_4b5, $offset_4b5);


    my $encryptedkey = substr($ciddata, $offset_4b5, 0x40);
    my $keydata="";
    for (my $i=0 ; $i<length($encryptedkey) ; $i+=32) {
        my $num = tonumber(substr($encryptedkey, $i, 32));
        my $decrypted = $num->copy()->bmodpow($dev->{exp}, $dev->{pq});
        my $decdata = todata($decrypted);

        $keydata.= $decdata;
        if (unpack("%8C*", substr($decdata, 0, 24)) != ord(substr($decdata, 30,1))) {
            printf("%s : %02x\n", unpack("H*", $decdata), unpack("%8C*", substr($decdata, 0, 24)));
            die "key-checksum mismatch\n";
        }
    }

    my $aeskey2= substr($keydata, 0, 16).substr($keydata, 32, 16);
    my $aesiv2= join("", map { substr($shadoc, $_, 1).substr($keydata, $_+16,1) } (0..7))
               .join("", map { substr($shadoc, $_+8, 1).substr($keydata, $_+48,1) } (0..7));

    #printf("aeskey2=%s\n", unpack("H*", $aeskey2));
    #printf("aesiv2=%s\n", unpack("H*", $aesiv2));

    my $ciddata_7ba= ord(substr($ciddata, 0x7ba, 1));
    my $offset_7ba= 0x1000+(($ciddata_7ba&0x1f)<<10);

    #printf("ciddata[7ba]= %02x -> offset %04x\n", $ciddata_7ba, $offset_7ba);
    my $encryptedlockdata= substr($ciddata, $offset_7ba, 0x800);
    #printf("enclockdata=%s\n", unpack("H*", $encryptedlockdata));

    my $lockdata= doaesdecrypt($encryptedlockdata, $aeskey2, $aesiv2);
    #printf("lockdata=%s\n", unpack("H*", $lockdata));

    if (swapdw(sha256(substr($lockdata, 0, 0x7e0))) ne substr($lockdata, 0x7e0, 0x20)) {
        printf("sha[000-7e0]=%s\n", unpack 'H*', swapdw(sha256(substr($lockdata, 0, 0x7e0))));
        warn "invalid data-checksum\n";
    }

    return $lockdata;
}

sub encrypt_cid_block {
    my ($ciddata, $lockdata, $docid)= @_;
    my $shadoc = swapdw(sha256($docid));
    #printf("encrypt_cid_block\n");

    # part2: decrypt the data section

    my $ciddata_4b5= ord(substr($ciddata, 0x4b5, 1));

    my $offset_4b5= 0xB000 + (($ciddata_4b5 &0xf)<<10);
    #printf("ciddata[4b5]= %02x -> offset %04x\n", $ciddata_4b5, $offset_4b5);


    my $encryptedkey = substr($ciddata, $offset_4b5, 0x40);
    my $keydata="";
    for (my $i=0 ; $i<length($encryptedkey) ; $i+=32) {
        my $num = tonumber(substr($encryptedkey, $i, 32));
        my $decrypted = $num->bmodpow($dev->{exp}, $dev->{pq});
        my $decdata = todata($decrypted);

        $keydata.= $decdata;
        if (unpack("%8C*", substr($decdata, 0, 24)) != ord(substr($decdata, 30,1))) {
            printf("%s : %02x\n", unpack("H*", $decdata), unpack("%8C*", substr($decdata, 0, 24)));
            die "key-checksum mismatch\n";
        }
    }

    my $aeskey2= substr($keydata, 0, 16).substr($keydata, 32, 16);
    my $aesiv2= join("", map { substr($shadoc, $_, 1).substr($keydata, $_+16,1) } (0..7))
               .join("", map { substr($shadoc, $_+8, 1).substr($keydata, $_+48,1) } (0..7));

    #printf("aeskey2=%s\n", unpack("H*", $aeskey2));
    #printf("aesiv2=%s\n", unpack("H*", $aesiv2));
    substr($lockdata, 0x7e0, 0x20)= swapdw(sha256(substr($lockdata, 0, 0x7e0)));
    #printf("lockdata=%s\n", unpack("H*", $lockdata));

    my $encryptedlockdata= doaesencrypt($lockdata, $aeskey2, $aesiv2);
    #printf("enclockdata=%s\n", unpack("H*", $encryptedlockdata));

    my $ciddata_7ba= ord(substr($ciddata, 0x7ba, 1));
    my $offset_7ba= 0x1000+(($ciddata_7ba&0x1f)<<10);

    #printf("ciddata[7ba]= %02x -> offset %04x\n", $ciddata_7ba, $offset_7ba);

    substr($ciddata, $offset_7ba, 0x800)= $encryptedlockdata;

    # part1: verify cid block checksum

    # these are used to calculate the aes keys
    my $sha64k = swapdw(sha256(substr($ciddata, 0, -0x20)));

    my $aeskey1= substr($shadoc, 0, 16) . substr($sha64k, 0, 16);
    my $aesiv1= substr($sha64k, 16, 16) . substr($shadoc, 0, 16);

    my $bufend= doaesencrypt($sha64k, $aeskey1, $aesiv1);
    substr($ciddata, -0x20, 0x20)= $bufend;

    return $ciddata;
}

sub hexdump {
    return join("\n", map { my $x=substr($_[0], $_*16, 16);$x =~ tr/[\x00-\x1f\x7f-\xff]/./; sprintf("%04x: %s  %s", $_*16, unpack("H*", substr($_[0], $_*16, 16)), $x); } (0..length($_[0])/16-1))."\n";
}

exit(0);
sub tonumber {
    my $data= shift;
    my $num= 0;
    for (reverse unpack("C*", $data)) {
        $num = $num*256 + $_;
    }
    return $num;
}
sub todata {
    my $num= shift;
    my @data;
    while ($num) {
        push @data, $num % 256;
        $num = $num / 256;
    }
    return pack("C32", @data);
}
sub doaesdecrypt {
    my ($in, $key, $iv)= @_;

    my $aes= Crypt::Rijndael->new($key);
    return ivdecrypt($aes, $iv, $in);
}
sub ivdecrypt {
    my ($aes, $iv, $in)= @_;
    #printf("ivdecrypt\n");
    my $out= "";
    my $iv_lo= substr($iv, 0, 16);
    my $iv_hi= substr($iv, 16, 16);
    for (my $i=0 ; $i < length($in) ; $i += 32)
    {
        my $crp_lo = substr($in, $i, 16);
        my $dat_lo= $aes->decrypt($crp_lo);
        $out .= $dat_lo ^ $iv_lo;
        $iv_lo = $crp_lo;
        #printf("%dlo: d=%s c=%s i=%s\n", $i/32, map { unpack("H*", $_) } ($dat_lo, $crp_lo, $iv_lo));

        my $crp_hi = substr($in, $i+16, 16);
        my $dat_hi= $aes->decrypt($crp_hi);
        $out .= $dat_hi ^ $iv_hi;
        $iv_hi = $crp_hi;
        #printf("%dhi: d=%s c=%s i=%s\n", $i/32, map { unpack("H*", $_) } ($dat_hi, $crp_hi, $iv_hi));
    }
    # d0 = aes(c0)    | c0 = unaes(d0)
    # o0= d0^iv0      | d0 = o0^iv0
    # iv1 = d0        .
    # d1 = aes(c1)
    # o1 .= d1^iv1
    # iv2 = d1
    #
    return $out;
}
sub doaesencrypt {
    my ($in, $key, $iv)= @_;

    my $aes= Crypt::Rijndael->new($key);
    return ivencrypt($aes, $iv, $in);
}
sub ivencrypt {
    my ($aes, $iv, $in)= @_;
    #printf("ivencrypt\n");
    my $out= "";
    my $iv_lo= substr($iv, 0, 16);
    my $iv_hi= substr($iv, 16, 16);
    for (my $i=0 ; $i < length($in) ; $i += 32)
    {
        my $dat_lo = substr($in, $i, 16) ^ $iv_lo;
        my $crp_lo= $aes->encrypt($dat_lo);
        $out .= $crp_lo;
        $iv_lo = $crp_lo;
        #printf("%dlo: d=%s c=%s i=%s\n", $i/32, map { unpack("H*", $_) } ($dat_lo, $crp_lo, $iv_lo));

        my $dat_hi = substr($in, $i+16, 16) ^ $iv_hi;
        my $crp_hi= $aes->encrypt($dat_hi);
        $out .= $crp_hi;
        $iv_hi = $crp_hi;
        #printf("%dhi: d=%s c=%s i=%s\n", $i/32, map { unpack("H*", $_) } ($dat_hi, $crp_hi, $iv_hi));
    }
    # d0 = o0^iv0
    # c0 = unaes(d0)
    # iv1 = d0
    return $out;
}

sub swapdw {
    return pack 'V*', unpack 'N*', $_[0];
}