package Crypt::JWT;

use strict;
use warnings;

our $VERSION = '0.002';

use Exporter 'import';
our %EXPORT_TAGS = ( all => [qw(decode_jwt encode_jwt)] );
our @EXPORT_OK = ( @{ $EXPORT_TAGS{'all'} } );
our @EXPORT = qw();

use Carp;
use MIME::Base64 qw(decode_base64url encode_base64url);
use JSON qw(decode_json encode_json);
use Crypt::PK::RSA;
use Crypt::PK::ECC;
use Crypt::PRNG qw(random_bytes);
use Crypt::KeyWrap ':all';
use Crypt::AuthEnc::GCM qw(gcm_encrypt_authenticate gcm_decrypt_verify);
use Crypt::Mac::HMAC qw(hmac);
use Compress::Raw::Zlib;

# JWS: https://tools.ietf.org/html/rfc7515
# JWE: https://tools.ietf.org/html/rfc7516
# JWK: https://tools.ietf.org/html/rfc7517
# JWA: https://tools.ietf.org/html/rfc7518
# JWT: https://tools.ietf.org/html/rfc7519

sub _prepare_rsa_key {
  my ($key) = @_;
  # we need Crypt::PK::RSA object
  return $key                       if ref($key) eq 'Crypt::PK::RSA';
  return Crypt::PK::RSA->new($key)  if ref($key) eq 'HASH';
  return Crypt::PK::RSA->new(@$key) if ref($key) eq 'ARRAY';
  return Crypt::PK::RSA->new(\$key) if !ref($key);
  # handle also: Crypt::OpenSSL::RSA, Crypt::X509, Crypt::OpenSSL::X509
  my $str;
  if (ref($key) eq 'Crypt::OpenSSL::RSA') {
    # https://metacpan.org/pod/Crypt::OpenSSL::RSA
    $str = $key->is_private ? $key->get_private_key_string : $key->get_public_key_string;
  }
  elsif (ref($key) =~ /^Crypt::(X509|OpenSSL::X509)$/) {
    # https://metacpan.org/pod/Crypt::X509
    # https://metacpan.org/pod/Crypt::OpenSSL::X509
    $str = $key->pubkey;
  }
  return Crypt::PK::RSA->new(\$str) if !ref($str);
  return undef;
}

sub _prepare_ecc_key {
  my ($key) = @_;
  # we need Crypt::PK::ECC object
  return $key                       if ref($key) eq 'Crypt::PK::ECC';
  return Crypt::PK::ECC->new($key)  if ref($key) eq 'HASH';
  return Crypt::PK::ECC->new(@$key) if ref($key) eq 'ARRAY';
  return Crypt::PK::ECC->new(\$key) if !ref($key);
  return undef;
}

sub _encrypt_cek {
  my ($key, $hdr) = @_;
  my $alg = $hdr->{alg};
  my $enc = $hdr->{enc};
  
  if ($alg eq 'dir') {
    return ($key, '');
  }
  
  my $cek;
  my $ecek;
  if ($enc =~ /^A(128|192|256)GCM/) {
    $cek = random_bytes($1/8);
  }
  elsif ($enc =~ /^A(128|192|256)CBC/) {
    $cek = random_bytes(2*$1/8);
  }

  if ($alg =~ /^A(128|192|256)KW$/) {
    $ecek = aes_key_wrap($key, $cek);
    return ($cek, $ecek);
  }
  elsif ($alg =~ /^A(128|192|256)GCMKW$/) {
    my ($t, $i);
    ($ecek, $t, $i) = gcm_key_wrap($key, $cek);
    $hdr->{tag} = encode_base64url($t);
    $hdr->{iv}  = encode_base64url($i);
    return ($cek, $ecek);
  }
  elsif ($alg =~ /^PBES2-HS(512|384|256)\+A(128|192|256)KW$/) {
    my $salt = random_bytes(16); #XXX
    my $iter = 5000;             #XXX
    $ecek = pbes2_key_wrap($key, $cek, $alg, $salt, $iter);
    $hdr->{p2s} = encode_base64url($salt);
    $hdr->{p2c} = $iter;
    return ($cek, $ecek);
  }
  elsif ($alg =~ /^RSA(-OAEP|-OAEP-256|1_5)$/) {
    $key = _prepare_rsa_key($key);
    $ecek = rsa_key_wrap($key, $cek, $alg);
    return ($cek, $ecek);
  }
  elsif ($alg =~ /^ECDH-ES\+A(128|192|256)KW$/) {
    $key = _prepare_ecc_key($key);
    ($ecek, $hdr->{epk}) = ecdhaes_key_wrap($key, $cek, $alg, $hdr->{apu}, $hdr->{apv});
    return ($cek, $ecek);
  }
  elsif ($alg eq 'ECDH-ES') {
    $key = _prepare_ecc_key($key);
    ($cek, $hdr->{epk}) = ecdh_key_wrap($key, $enc, $hdr->{apu}, $hdr->{apv});
    return ($cek, '');
  }
  die "XXX";
}

sub _decrypt_cek {
  my ($ecek, $key, $hdr) = @_;
  my $alg = $hdr->{alg};
  my $enc = $hdr->{enc};
  if ($alg eq 'dir') {
    return $key;
  }
  elsif ($alg =~ /^A(128|192|256)KW$/) {
    return aes_key_unwrap($key, $ecek);
  }
  elsif ($alg =~ /^A(128|192|256)GCMKW$/) {
    return gcm_key_unwrap($key, $ecek, decode_base64url($hdr->{tag}), decode_base64url($hdr->{iv}));
  }
  elsif ($alg =~ /^PBES2-HS(512|384|256)\+A(128|192|256)KW$/) {
    return pbes2_key_unwrap($key, $ecek, $alg, decode_base64url($hdr->{p2s}), $hdr->{p2c});
  }
  elsif ($alg =~ /^RSA(-OAEP|-OAEP-256|1_5)$/) {
    $key = _prepare_rsa_key($key);
    return rsa_key_unwrap($key, $ecek, $alg);
  }
  elsif ($alg =~ /^ECDH-ES\+A(128|192|256)KW$/) {
    $key = _prepare_ecc_key($key);
    return ecdhaes_key_unwrap($key, $ecek, $alg, $hdr->{epk}, $hdr->{apu}, $hdr->{apv});
  }
  elsif ($alg eq 'ECDH-ES') {
    $key = _prepare_ecc_key($key);
    return ecdh_key_unwrap($key, $enc, $hdr->{epk}, $hdr->{apu}, $hdr->{apv});
  }
  die "XXX";
}

sub _encrypt_payload {
  my ($cek, $enc, $encoded_header, $payload) = @_;
  if ($enc =~ /^A(128|192|256)GCM$/) {
    # https://tools.ietf.org/html/rfc7518#section-5.3
    my $len1 = $1/8;
    my $len2 = length($cek);
    die "jwe: wrong AES key length ($len1 vs. $len2) for $enc" unless $len1 == $len2;
    my $iv = random_bytes(16); # for AES always 16
    my ($ct, $tag) = gcm_encrypt_authenticate('AES', $cek, $iv, $encoded_header, $payload);
    return ($ct, $iv, $tag);
  }
  elsif ($enc =~ /^A(128|192|256)CBC-HS(256|384|512)$/) {
    # https://tools.ietf.org/html/rfc7518#section-5.2
    my ($size, $hash) = ($1/8, "SHA$2");
    my $key_len = length($cek) / 2;
    my $mac_key = substr($cek, 0, $key_len);
    my $aes_key = substr($cek, $key_len, $key_len);
    die "jwe: wrong AES key length ($key_len vs. $size)" unless $key_len == $size;
    my $iv = random_bytes(16); # for AES always 16
    my $m = Crypt::Mode::CBC->new('AES');
    my $ct = $m->encrypt($payload, $aes_key, $iv);
    my $aad_len = length($encoded_header); # AAD == original encoded header
    my $mac_input = $encoded_header . $iv . $ct . pack('N2', ($aad_len / 2147483647)*8, ($aad_len % 2147483647)*8);
    my $mac = hmac($hash, $mac_key, $mac_input);
    my $sig_len = length($mac) / 2;
    my $sig = substr($mac, 0, $sig_len);
    return ($ct, $iv, $sig);
  }  
  die "jwe: unsupported enc '$enc'";
}

sub _decrypt_payload {
  my ($cek, $enc, $encoded_header, $ct, $iv, $tag) = @_;
  if ($enc =~ /^A(128|192|256)GCM$/) {
    # https://tools.ietf.org/html/rfc7518#section-5.3
    my $len1 = $1/8;
    my $len2 = length($cek);
    die "jwe: wrong AES key length ($len1 vs. $len2) for $enc" unless $len1 == $len2;
    return gcm_decrypt_verify('AES', $cek, $iv, $encoded_header, $ct, $tag);
  }
  elsif ($enc =~ /^A(128|192|256)CBC-HS(256|384|512)$/) {
    # https://tools.ietf.org/html/rfc7518#section-5.2
    my ($size, $hash) = ($1/8, "SHA$2");
    my $key_len = length($cek) / 2;
    my $mac_key = substr($cek, 0, $key_len);
    my $aes_key = substr($cek, $key_len, $key_len);
    die "jwe: wrong AES key length ($key_len vs. $size)" unless $key_len == $size;
    my $aad_len = length($encoded_header); # AAD == original encoded header
    my $mac_input = $encoded_header . $iv . $ct . pack('N2', ($aad_len / 2147483647)*8, ($aad_len % 2147483647)*8);
    my $mac = hmac($hash, $mac_key, $mac_input);
    my $sig_len = length($mac) / 2;
    my $sig = substr($mac, 0, $sig_len);
    die "jwe: tag mismatch" unless $sig eq $tag;
    my $m = Crypt::Mode::CBC->new('AES');
    my $pt = $m->decrypt($ct, $aes_key, $iv);
    return $pt;
  }
  die "jwe: unsupported enc '$enc'";
}

sub _decode_jwe {
  my ($header, $ecek, $ct, $iv, $aad, $tag, %args) = @_;
  my $key = defined $args{keypass} ? [$args{key}, $args{keypass}] : $args{key};
  my $cek = _decrypt_cek($ecek, $key, $header);
  my $pt  = _decrypt_payload($cek, $header->{enc}, $aad, $ct, $iv, $tag);
  if (my $zip = $header->{zip}) {
    if ($zip eq "DEF") {
      my $d = Compress::Raw::Zlib::Inflate->new(-Bufsize => 300, -WindowBits => -MAX_WBITS());
      my $output = '';
      $d->inflate($pt, $output);
      die "JWT: inflate failed" unless $output;
      return $output;
    }
    die "JWT: unknown zip method '$zip'";
  }
  return $pt;
}

sub _sign_jws {
  my ($header, $b64u_payload, $b64u_header, %args) = @_;
  my $key = $args{key};
  my $pass = $args{keypass};
  my $alg = $header->{alg};
  if ($alg eq 'none' ) { # no integrity
    return '';
  }
  my $sig;
  my $data = "$b64u_header.$b64u_payload";
  if ($alg =~ /^HS(256|384|512)$/) { # HMAC integrity
    $sig = hmac("SHA$1", $key, $data);
  }
  elsif ($alg =~ /^RS(256|384|512)/) { # RSA+PKCS1-V1_5 signatures
    my $pk = _prepare_rsa_key($key, $pass);
    $sig  = $pk->sign_message($data, "SHA$1", 'v1.5');
  }
  elsif ($alg =~ /^PS(256|384|512)/) { # RSA+PSS signatures
    my $hash = "SHA$1";
    my $hashlen = $1/8;
    my $pk = _prepare_rsa_key($key, $pass);
    $sig  = $pk->sign_message($data, $hash, 'pss', $hashlen);
  }
  elsif ($alg =~ /^ES(256|384|512)/) { # ECDSA signatures
    my $pk = _prepare_ecc_key($key, $pass);
    $sig  = $pk->sign_message_rfc7518($data, "SHA$1");
  }
  return encode_base64url($sig);
}

sub _verify_jws {
  my ($header, $b64u_payload, $b64u_header, $b64u_sig, %args) = @_;
  my $key = $args{key};
  my $pass = $args{keypass};
  my $alg = $header->{alg};
  die "JWS: alg 'none' not allowed" if $alg eq 'none' && !$args{allow_none};
 
  my $data = "$b64u_header.$b64u_payload";
  my $sig = decode_base64url($b64u_sig);

  if ($alg eq 'none' ) { # no integrity
    return 1;
  }
  elsif ($alg =~ /^HS(256|384|512)$/) { # HMAC integrity
    return 1 if $sig eq hmac("SHA$1", $key, $data);
  }
  elsif ($alg =~ /^RS(256|384|512)/) { # RSA+PKCS1-V1_5 signatures
    my $hash = "SHA$1";
    my $pk = _prepare_rsa_key($key);
    return 1 if $pk->verify_message($sig, $data, $hash, 'v1.5');
  }
  elsif ($alg =~ /^PS(256|384|512)/) { # RSA+PSS signatures
    my $hash = "SHA$1";
    my $hashlen = $1/8;
    my $pk = _prepare_rsa_key($key);
    return 1 if  $pk->verify_message($sig, $data, $hash, 'pss', $hashlen);
  }
  elsif ($alg =~ /^ES(256|384|512)/) { # ECDSA signatures
    my $hash = "SHA$1";
    my $pk = _prepare_ecc_key($key);
    return 1 if $pk->verify_message_rfc7518($sig, $data, $hash);
  }
  return 0;
}

sub decode_jwt {
  my %args = @_;
  my $token = $args{token};
  return undef unless defined $token;

  if ($token =~ /^([a-zA-Z0-9_-]+)\.([a-zA-Z0-9_-]*)\.([a-zA-Z0-9_-]+)\.([a-zA-Z0-9_-]+)\.([a-zA-Z0-9_-]+)$/) {
    # JWE token (5 segments)
    my $encoded_header = $1;
    my $header = decode_base64url($1);
    my $ecek   = decode_base64url($2);
    my $iv     = decode_base64url($3);
    my $ct     = decode_base64url($4);
    my $tag    = decode_base64url($5);
    utf8::decode($header) if $header && !utf8::is_utf8($header);
    return undef unless defined $header && defined $ecek && defined $iv && defined $ct && defined $tag;
    $header = decode_json($header);
    return undef unless defined $header;
    my $payload = _decode_jwe($header, $ecek, $ct, $iv, $encoded_header, $tag, %args);
    return ($header, $payload) if $args{decode_header};
    return $payload;
  }
  elsif ($token =~ /^([a-zA-Z0-9_-]+)\.([a-zA-Z0-9_-]+)\.([a-zA-Z0-9_-]*)$/) {
    # JWS token (3 segments)
    my $b64u_header  = $1;
    my $b64u_payload = $2;
    my $b64u_sig     = $3;
    my $header  = decode_base64url($b64u_header);
    utf8::decode($header) if $header && !utf8::is_utf8($header);
    return undef unless defined $header;
    $header = decode_json($header);
    return undef unless defined $header;
    die "JWS: invalid signature" unless _verify_jws($header, $b64u_payload, $b64u_header, $b64u_sig, %args);
    my $payload = decode_base64url($b64u_payload);
    return ($header, $payload) if $args{decode_header};
    return $payload;
  }
  die "decode_jwt: invalid token format";
}

sub encode_jwt {
  my %args = @_;
  my $key     = $args{key};
  my $payload = $args{payload};
  my $alg     = $args{alg};
  my $enc     = $args{enc};
  my $header  = $args{extra_headers} ? \%{$args{extra_headers}} : {};

  if ($alg =~ /^(none|((HS|RS|PS|ES)(512|384|256)))$/) {
    ###JWS
    # prepare header
    $header->{iat} = time if $args{auto_iat};
    $header->{alg} = $alg;
    die "JWS: alg 'none' not allowed" if $alg eq 'none' && !$args{allow_none};
    # encode header
    my $json_header = encode_json($header);
    my $b64u_header = encode_base64url($json_header);
    # encode payload
    my $data = ref($payload) =~ /^(HASH|ARRAY)$/ ? encode_json($payload) : $payload;
    my $b64u_payload = encode_base64url($data);
    # sign header
    my $b64u_signature = _sign_jws($header, $b64u_payload, $b64u_header, %args);
    # create token
    my $token = "$b64u_header.$b64u_payload.$b64u_signature";
    return $token;
  }
  else {
    ### JWT
    # prepare header
    $header->{iat} = time if 0 && 'XXX-FIXME';
    $header->{alg} = $alg;
    $header->{enc} = $enc;
    # prepare cek
    my ($cek, $ecek) = _encrypt_cek($key, $header);
    # encode header
    my $json_header = encode_json($header);
    my $b64u_header = encode_base64url($json_header);
    # compress payload
    #XXX
    # encrypt payload
    my ($ct, $iv, $tag) = _encrypt_payload($cek, $enc, $b64u_header, $payload);
    # create token
    my $token = $b64u_header . "." .
                encode_base64url($ecek) . "." .
                encode_base64url($iv) . "." .
                encode_base64url($ct) . "." .
                encode_base64url($tag);
    return $token;
  }
  die "XXX alg=$alg";
}


1;

# https://metacpan.org/pod/JSON::WebToken
# https://metacpan.org/pod/Mojo::JWT
# https://bitbucket.org/b_c/jose4j/wiki/JWE%20Examples
# https://bitbucket.org/b_c/jose4j/wiki/JWS%20Examples

#XXX
# https://github.com/dvsekhvalnov/jose-jwt/tree/master/JWT/jwe
### encode:
### - croak => 0|1
### - payload
### - alg
### - enc
### - key
### - keypass
### - extra_headers
### - allow_none
### - auto_iat
### - ref2json
### ???auto_jti => 0|<len>
### ???deflate => 3
### ???sereal  => 3
### 
### decode:
### - croak => 0|1
### - token
### - alg
### - enc
### - key
### - keypass
### ???keystash
### ???leeway = > 30s OR verify_iat=>-3
### - accepted_alg
### - accepted_enc
### - verify_iat        undef|0|1|sub
### - verify_nbf        undef|0|1|sub
### - verify_exp        undef|0|1|sub
### - verify_iss        undef|sub|regexp
### - verify_aud        undef|sub|regexp
### - verify_sub        undef|sub|regexp
### - verify_jti        undef|sub|regexp
### - verify_signature  0|1
### - json2ref
### 
### https://github.com/progrium/ruby-jwt
### verify_iss
### verify_aud
### verify_jti
### verify_sub
### (!! jti = hmac(iat
### 
### leeway = > 30s
### allow_none
### allow_bad_nbf
### allow_bad_exp
### allow_bad_signature
### 
### key
### keystash
### keypass
### 
### https://github.com/jpadilla/pyjwt/
###    'verify_signature': True,
###    'verify_exp': True,  undef=no, 0=presne, 1=tolerance+-1
###    'verify_nbf': True,  undef=no, 0=presne, 1=tolerance+-1
###    'verify_iat': True,  undef=no, 0=presne, 1=tolerance+-1
###    'verify_aud': True
###    'require_exp': False,
###    'require_iat': False,
###    'require_nbf': False
   
   
=pod

=head1 NAME

Crypt::JWT - JSON Web Token (JWT, JWS, JWE) as defined by RFC7519, RFC7515, RFC7516

=head1 SYNOPSIS

   # encoding
   use Crypt::JWT qw(encode_jwt);
   my $jws_token = encode_jwt(payload=>$data, alg=>'HS256', key=>'secret'); 
   my $jwe_token = encode_jwt(payload=>$data, alg=>'A192GCMKW', enc=>'A192CBC_HS384', key=>'secret'); 

   # decoding
   use Crypt::JWT qw(decode_jwt);
   my $data1 = decode_jwt(token=>$jws_token, key=>'secret'); 
   my $data2 = decode_jwt(token=>$jwe_token, key=>'secret');

=head1 DESCRIPTION

B<BEWARE:> experimental, unfinished, unstable, work in progress!!!

Implements B<JSON Web Token (JWT)> - L<https://tools.ietf.org/html/rfc7519>.
The implementation covers not only B<JSON Web Signature (JWS)> - L<https://tools.ietf.org/html/rfc7515>,
but also B<JSON Web Encryption (JWE)> - L<https://tools.ietf.org/html/rfc7516>.

The module implements B<all (100%) algorithms> defined in L<https://tools.ietf.org/html/rfc7518> - B<JSON Web Algorithms (JWA)>.

Supported JWE (encryption) algorithms:

 alg                 enc
 ------------------  -------------
 dir                 A128GCM
 A128KW              A192GCM
 A192KW              A256GCM
 A256KW              A128CBC-HS256
 A128GCMKW           A192CBC-HS384
 A192GCMKW           A256CBC-HS512
 A256GCMKW           
 PBES2-HS256+A128KW  
 PBES2-HS384+A192KW  
 PBES2-HS512+A256KW  
 RSA-OAEP            
 RSA-OAEP-256        
 RSA1_5              
 ECDH-ES+A128KW
 ECDH-ES+A192KW
 ECDH-ES+A256KW
 ECDH-ES

Supported JWS (signature) algorithms:

 alg        note
 ---------  ----------------------------------------
 none       no integrity (NOTE: disabled by default)
 HS256      HMAC integrity
 HS384
 HS512
 RS256      RSA+PKCS1-V1_5 signatures
 RS384
 RS512
 PS256      RSA+PSS signatures
 PS384
 PS512
 ES256      ECDSA signatures
 ES384
 ES512

=head1 FUNCTIONS

=head2 decode_jwt

 my $data = decode_jwt(%named_args);
 
Named arguments:

=over

=item token

XXX-TODO

=item key

XXX-TODO

=item allow_none

XXX-TODO

=back

=head2 encode_jwt

 my $token = encode_jwt(%named_args);

Named arguments:

=over

=item payload

XXX-TODO

=item key

XXX-TODO

=item allow_none

XXX-TODO

=item auto_iat

XXX-TODO

=back

=head1 SEE ALSO

L<Crypt::Cipher::AES>, L<Crypt::AuthEnc::GCM>, L<Crypt::PK::RSA>, L<Crypt::PK::ECC>, L<Crypt::KeyDerivation>, L<Crypt::KeyWrap>
