% Licensed under the Apache License, Version 2.0 (the "License"); you may not
% use this file except in compliance with the License. You may obtain a copy of
% the License at
%
%   http://www.apache.org/licenses/LICENSE-2.0
%
% Unless required by applicable law or agreed to in writing, software
% distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
% WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
% License for the specific language governing permissions and limitations under
% the License.

-module(aegis_keywrap).
-include("aegis.hrl").

%% Implementation of NIST Special Publication 800-38F
%% For wrapping and unwrapping keys with AES.

-export([key_wrap/2, key_unwrap/2]).

-define(ICV1, 16#A6A6A6A6A6A6A6A6).

-spec key_wrap(WrappingKey :: aegis_key_manager:key_fun(), KeyToWrap :: aegis_key_manager:key_fun()) -> binary().
key_wrap(WrappingKeyFun, KeyToWrapFun)
  when is_function(WrappingKeyFun, 0), is_function(KeyToWrapFun, 0) ->
    N = bit_size(KeyToWrapFun()) div 64,
    wrap(WrappingKeyFun, <<?ICV1:64>>, KeyToWrapFun(), 1, 6 * N).

wrap(_WrappingKeyFun, A, R, T, End) when T > End ->
    <<A/binary, R/binary>>;
wrap(WrappingKeyFun, A, R, T, End) ->
    <<R1:64, Rest/binary>> = R,
    <<MSB_B:64, LSB_B:64>> = ?aes_ecb_encrypt(WrappingKeyFun(), <<A/binary, R1:64>>),
    wrap(WrappingKeyFun, <<(MSB_B bxor T):64>>, <<Rest/binary, LSB_B:64>>, T + 1, End).


-spec key_unwrap(WrappingKey :: aegis_key_manager:key_fun(), KeyToUnwrap :: binary()) -> aegis_key_manager:key_fun() | fail.
key_unwrap(WrappingKeyFun, KeyToUnwrap)
  when is_function(WrappingKeyFun, 0), bit_size(KeyToUnwrap) rem 64 == 0 ->
    N = (bit_size(KeyToUnwrap) div 64),
    <<A:64, R/binary>> = KeyToUnwrap,
    case unwrap(WrappingKeyFun, <<A:64>>, R, 6 * (N - 1)) of
        <<?ICV1:64, UnwrappedKey/binary>> ->
            fun() -> UnwrappedKey end;
        _ ->
            fail
    end.

unwrap(_WrappingKeyFun, A, R, 0) ->
    <<A/binary, R/binary>>;
unwrap(WrappingKeyFun, <<A:64>>, R, T) ->
    RestSize = bit_size(R) - 64,
    <<Rest:RestSize, R2: 64>> = R,
    <<MSB_B:64, LSB_B:64>> = ?aes_ecb_decrypt(WrappingKeyFun(), <<(A bxor T):64, R2:64>>),
    unwrap(WrappingKeyFun, <<MSB_B:64>>, <<LSB_B:64, Rest:RestSize>>, T - 1).


-ifdef(TEST).
-include_lib("eunit/include/eunit.hrl").

wrap_test_() ->
    [
     %% 128 KEK / 128 DATA
     test_wrap_unwrap(<<16#000102030405060708090A0B0C0D0E0F:128>>,
                      <<16#00112233445566778899AABBCCDDEEFF:128>>,
                      <<16#1FA68B0A8112B447AEF34BD8FB5A7B829D3E862371D2CFE5:192>>),
     %% 192 KEK / 128 DATA
     test_wrap_unwrap(<<16#000102030405060708090A0B0C0D0E0F1011121314151617:192>>,
                      <<16#00112233445566778899AABBCCDDEEFF:128>>,
                      <<16#96778B25AE6CA435F92B5B97C050AED2468AB8A17AD84E5D:192>>),
     %% 256 KEK / 128 DATA
     test_wrap_unwrap(<<16#000102030405060708090A0B0C0D0E0F101112131415161718191A1B1C1D1E1F:256>>,
                      <<16#00112233445566778899AABBCCDDEEFF:128>>,
                      <<16#64E8C3F9CE0F5BA263E9777905818A2A93C8191E7D6E8AE7:192>>),
     %% 192 KEK / 192 DATA
     test_wrap_unwrap(<<16#000102030405060708090A0B0C0D0E0F1011121314151617:192>>,
                      <<16#00112233445566778899AABBCCDDEEFF0001020304050607:192>>,
                      <<16#031D33264E15D33268F24EC260743EDCE1C6C7DDEE725A936BA814915C6762D2:256>>),
     %% 256 KEK / 192 DATA
     test_wrap_unwrap(<<16#000102030405060708090A0B0C0D0E0F101112131415161718191A1B1C1D1E1F:256>>,
                      <<16#00112233445566778899AABBCCDDEEFF0001020304050607:192>>,
                      <<16#A8F9BC1612C68B3FF6E6F4FBE30E71E4769C8B80A32CB8958CD5D17D6B254DA1:256>>),
     %% 256 KEK / 256 DATA
     test_wrap_unwrap(<<16#000102030405060708090A0B0C0D0E0F101112131415161718191A1B1C1D1E1F:256>>,
                      <<16#00112233445566778899AABBCCDDEEFF000102030405060708090A0B0C0D0E0F:256>>,
                      <<16#28C9F404C4B810F4CBCCB35CFB87F8263F5786E2D80ED326CBC7F0E71A99F43BFB988B9B7A02DD21:320>>)].

test_wrap_unwrap(WrappingKey, KeyToWrap, ExpectedWrappedKey) ->
    [?_assertEqual(ExpectedWrappedKey, key_wrap(WrappingKey, KeyToWrap)),
     ?_assertEqual(KeyToWrap, key_unwrap(WrappingKey, key_wrap(WrappingKey, KeyToWrap)))].

fail_test() ->
    KEK = <<16#000102030405060708090A0B0C0D0E0F101112131415161718191A1B1C1D1E1F:256>>,
    CipherText = <<16#28C9F404C4B810F4CBCCB35CFB87F8263F5786E2D80ED326CBC7F0E71A99F43BFB988B9B7A02DD20:320>>,
    ?assertEqual(fail, key_unwrap(KEK, CipherText)).

-endif.