%%% Copyright (c) 2014, NORDUnet A/S.
%%% See LICENSE for licensing information.
%%%
%%% @doc Server holding log entries in a database and hashes in a
%%% Merkle tree. A backend for things like Certificate Transparency
%%% (RFC 6962).
%%%
%%% When you submit data for insertion in the log, it's stored in an
%%% append only database with an accompanying Merkle tree. The leaves
%%% of the tree hold hashes of submitted data and makes it possible
%%% for anyone to verify wether a given piece of data is or is not
%%% present in the log.
%%%
%%% In return you will get a signed timestamp which is a promise that
%%% your entry will be present in the log within a certain time period
%%% (the MMD). This signed timestamp can later, together with the
%%% public key of the log, be used to ensure that your entry is indeed
%%% present in the log.

%%% TODO
%%% - get rid of CT-specific stuff that has creeped in

-module(plop).
-behaviour(gen_server).

%% API.
-export([start_link/2, stop/0]).
-export([get_logid/0, serialise/1]).
-export([add/1, sth/0, get/2, consistency/2, inclusion/2]).
%% API for tests.
-export([read_keyfile_rsa/2, read_keyfiles_ec/2]).
-export([testing_get_pubkey/0]).
%% gen_server callbacks.
-export([init/1, handle_call/3, terminate/2,
         handle_cast/2, handle_info/2, code_change/3]).

-include("$CTROOT/plop/include/plop.hrl").
-include("db.hrl").
-include_lib("public_key/include/public_key.hrl").
-include_lib("eunit/include/eunit.hrl").

-define(TESTPRIVKEYFILE, "test/eckey.pem").
-define(TESTPUBKEYFILE, "test/eckey-public.pem").

-record(state, {pubkey :: public_key:rsa_public_key(),
                privkey :: public_key:rsa_private_key(),
                logid :: binary()}).

%% @doc The parts of an STH which is to be signed. Used as the
%% interface to plop:sth/1, for testing.
-record(sth_signed, {
          version = ?PLOPVERSION :: non_neg_integer(),
          signature_type :: signature_type(),
          timestamp = now :: 'now' | integer(),
          tree_size :: integer(),
          root_hash :: binary()                 % SHA-256
         }).
-type sth_signed() :: #sth_signed{}.

%% @doc What's signed in an SPT. Used for serialisation before hasning
%% and signing. FIXME: Overlapping #spt{} -- merge somehow?
-record(spt_signed, {
          version = ?PLOPVERSION :: non_neg_integer(),
          signature_type :: signature_type(),
          timestamp :: integer(),
          entry_type :: entry_type(),
          signed_entry :: binary()
         }).
-type spt_signed() :: #spt_signed{}.

start_link(Keyfile, Passphrase) ->
    gen_server:start_link({local, ?MODULE}, ?MODULE, [Keyfile, Passphrase], []).

stop() ->
    gen_server:call(?MODULE, stop).

%%%%%%%%%%%%%%%%%%%%
init([PrivKeyfile, PubKeyfile]) ->
    %% Read RSA keypair.
    %% {Private_key, Public_key} = read_keyfile_rsa(Keyfile, Passphrase),
    %% LogID = crypto:hash(sha256,
    %%                     public_key:der_encode('RSAPublicKey', Public_key)),
    %% Read EC keypair.
    {Private_key, Public_key} = read_keyfiles_ec(PrivKeyfile, PubKeyfile),
    LogID = crypto:hash(sha256, public_key:der_encode(
                                  'ECPoint',
                                  element(2, element(1, Public_key)))), % FIXME!
    _Tree = ht:reset_tree([db:size() - 1]),
    {ok, #state{pubkey = Public_key,
                privkey = Private_key,
                logid = LogID}}.

handle_cast(_Request, State) ->
    {noreply, State}.

handle_info(_Info, State) ->
    {noreply, State}.

code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

terminate(_Reason, _State) ->
    io:format("~p terminating~n", [?MODULE]),
    ok.

%%%%%%%%%%%%%%%%%%%%
add(Data) when is_record(Data, timestamped_entry) ->
    gen_server:call(?MODULE, {add, Data}).
sth() ->
    gen_server:call(?MODULE, {sth, []}).
get(Start, End) ->
    gen_server:call(?MODULE, {get, {Start, End}}).
consistency(TreeSizeFirst, TreeSizeSecond) ->
    gen_server:call(?MODULE, {consistency, {TreeSizeFirst, TreeSizeSecond}}).
inclusion(Hash, TreeSize) ->
    gen_server:call(?MODULE, {inclusion, {Hash, TreeSize}}).
get_logid() ->
    gen_server:call(?MODULE, {get, logid}).
testing_get_pubkey() ->
    gen_server:call(?MODULE, {test, pubkey}).
%%%%%%%%%%%%%%%%%%%%
handle_call(stop, _From, State) ->
    {stop, normal, stopped, State};

%% FIXME: What's the right interface for add()? Need to be able to set
%% version and signature type, at least. That's missing from
%% #timestamped_entry, so add it somehow.
handle_call({add, #timestamped_entry{
                     timestamp = Timestamp_in, entry = Entry}},
            _From,
            State = #state{privkey = Privkey,
                           logid = LogID}) ->
    TimestampedEntry = #timestamped_entry{
                          timestamp = timestamp(Timestamp_in),
                          entry = Entry},
    {ok, SPT} = do_add(TimestampedEntry, Privkey, LogID),
    {reply, SPT, State};

handle_call({sth, Data}, _From,
            Plop = #state{privkey = PrivKey}) ->
    {reply, sth(PrivKey, Data), Plop};

handle_call({get, {Start, End}}, _From, Plop) ->
    {reply, db:get_by_index(Start, End), Plop};

handle_call({get, logid}, _From,
            Plop = #state{logid = LogID}) ->
    {reply, LogID, Plop};

handle_call({consistency, {First, Second}}, _From, Plop) ->
    {reply, ht:consistency(First - 1, Second - 1), Plop};

handle_call({inclusion, {Hash, TreeSize}}, _From, Plop) ->
    {Index, Proof} = case db:find(Hash) of
                         [] -> [];
                         {plop, I, _Hash, _MTL, _SPT} ->
                             {I, ht:path(I, TreeSize - 1)}
            end,
    {reply, {Index, Proof}, Plop};

handle_call({test, pubkey}, _From,
            Plop = #state{pubkey = PK}) ->
    {reply, PK, Plop}.

%%%%%%%%%%%%%%%%%%%%
%% db_get_single_entry(N) ->
%%     [#mtl{entry = #timestamped_entry{entry = #plop_entry{data = Data}}}] =
%%         db:get_by_index(N, N),
%%     Data.

-spec do_add(timestamped_entry(),
             public_key:rsa_private_key(),
             binary()) -> {ok|error, binary()}.
do_add(TimestampedEntry = #timestamped_entry{entry = PlopEntry},
       Privkey,
       LogID) ->
    DB_hash = crypto:hash(sha256, serialise(PlopEntry)),
    Record = db:find(DB_hash),
    case Record of
        #plop{index = _I, mtl = #mtl{entry = E}, spt = SPT} ->
            %%io:format("Found entry: index=~p~n", [I]),
            %% Database consistency checking. FIXME: Remove.
            Record = Record#plop{
                       hash = DB_hash,
                       mtl = #mtl{entry =
                                      #timestamped_entry{
                                         timestamp = E#timestamped_entry.timestamp,
                                         entry = PlopEntry}}},
            {ok, SPT};
        [] ->
            NewSPT = spt(LogID, Privkey, TimestampedEntry),
            MTL = #mtl{entry = TimestampedEntry},
            %%io:format("Creating new entry: index=~p~n", [ht:size()]),
            DB_data = #plop{index = ht:size(),
                            hash = DB_hash,
                            mtl = MTL,
                            spt = NewSPT},
            {atomic, ok} = db:add(DB_data),
            {ht:add(serialise(MTL)), NewSPT};
        Err -> {error, Err}
    end.
    
%% @doc Signed Plop Timestamp, conformant to an SCT in RFC6962 3.2 and
%% RFC5246 4.7.
-spec spt(binary(), public_key:rsa_private_key(), timestamped_entry()) -> spt().
spt(LogID, PrivKey, #timestamped_entry{
                       timestamp = Timestamp,
                       entry = #plop_entry{type = EntryType, data = EntryData}
                      }) ->
    BinToSign = serialise(#spt_signed{
                             signature_type = certificate_timestamp,
                             timestamp = Timestamp,
                             entry_type = EntryType,
                             signed_entry = EntryData}),
    Signature = #signature{
                   algorithm = #sig_and_hash_alg{
                                  hash_alg = sha256,
                                  signature_alg = ecdsa},
                   signature = signhash(BinToSign, PrivKey)},
    #spt{
       version = ?PLOPVERSION,
       logid = LogID,
       timestamp = Timestamp,
       signature = Signature}.

%% @doc Signed Tree Head as specified in RFC6962 section 3.2.
-spec sth(#'ECPrivateKey'{}, sth_signed() | list()) -> sth().
sth(PrivKey, []) ->
    sth(PrivKey, #sth_signed{timestamp = now});
sth(PrivKey, #sth_signed{version = Version, timestamp = Timestamp_in}) ->
    Timestamp = timestamp(Timestamp_in),
    Treesize = ht:size(),
    Roothash = ht:tree_hash(),
    BinToSign = serialise(#sth_signed{
                             version = Version,
                             signature_type = tree_hash,
                             timestamp = Timestamp,
                             tree_size = Treesize,
                             root_hash = Roothash}),
    Signature = #signature{
                   algorithm = #sig_and_hash_alg{
                                  hash_alg = sha256,
                                  signature_alg = ecdsa},
                   signature = signhash(BinToSign, PrivKey)},
    STH = #sth{
             treesize = Treesize,
             timestamp = Timestamp,
             roothash = Roothash,
             signature = Signature},
    %%io:format("STH: ~p~nBinToSign: ~p~nSignature: ~p~nTimestamp: ~p~n",
    %%          [STH, BinToSign, Signature, Timestamp]),
    STH.

%% TODO: Merge the keyfile reading functions.
%% @doc Read one password protected PEM file with an RSA keypair.
read_keyfile_rsa(Filename, Passphrase) ->
    {ok, PemBin} = file:read_file(Filename),
    [KeyPem] = public_key:pem_decode(PemBin),   % Use first entry.
    Privatekey = decode_key(KeyPem, Passphrase),
    {Privatekey, public_key(Privatekey)}.

%% @doc Read two PEM files, one with a private EC key and one with the
%% corresponding public EC key.
read_keyfiles_ec(PrivkeyFile, Pubkeyfile) ->
    {ok, PemBinPriv} = file:read_file(PrivkeyFile),
    [OTPPubParamsPem, PrivkeyPem] = public_key:pem_decode(PemBinPriv),
    Privatekey = decode_key(PrivkeyPem),

    {_, ParamsBin, ParamsEnc} = OTPPubParamsPem,
    PubParamsPem = {'EcpkParameters', ParamsBin, ParamsEnc},
    Params = public_key:pem_entry_decode(PubParamsPem),

    {ok, PemBinPub} = file:read_file(Pubkeyfile),
    [SPKIPem] = public_key:pem_decode(PemBinPub),
    %% SPKI is missing #'AlgorithmIdentifier' so pem_entry_decode won't do.
    %% Publickey = public_key:pem_entry_decode(SPKIPem),
    #'SubjectPublicKeyInfo'{algorithm = AlgoDer} = SPKIPem,
    SPKI = public_key:der_decode('SubjectPublicKeyInfo', AlgoDer),
    #'SubjectPublicKeyInfo'{subjectPublicKey = {_, Octets}} = SPKI,
    Point = #'ECPoint'{point = Octets},
    Publickey = {Point, Params},

    {Privatekey, Publickey}.

decode_key(Entry) ->
    public_key:pem_entry_decode(Entry).
decode_key(Entry, Passphrase) ->
    public_key:pem_entry_decode(Entry, Passphrase).

public_key(#'RSAPrivateKey'{modulus = Mod, publicExponent = Exp}) ->
    #'RSAPublicKey'{modulus = Mod, publicExponent = Exp}.


%% FIXME: Merge RSA and DC.
signhash(D, K) ->
    signhash_ec(D, K).

-spec signhash_ec(iolist() | binary(), public_key:ec_private_key()) -> binary().
signhash_ec(Data, PrivKey) ->
    public_key:sign(Data, sha256, PrivKey).

%% -spec signhash_rsa(iolist() | binary(), public_key:rsa_private_key()) -> binary().
%% signhash_rsa(Data, PrivKey) ->
%%     %% Was going to just crypto:sign/3 the hash but looking at
%%     %% digitally_signed() in lib/ssl/src/ssl_handshake.erl it seems
%%     %% like we should rather use (undocumented) encrypt_private/3.
%%     %public_key:sign(hash(sha256, BinToSign), sha256, PrivKey)
%%     public_key:encrypt_private(crypto:hash(sha256, Data),
%%                                PrivKey,
%%                                [{rsa_pad, rsa_pkcs1_padding}]).

%%%%%%%%%%%%%%%%%%%%
%% Serialisation of data.

%% FIXME: Make the type conversion functions return binaries instead
%% of integers. Moving knowledge about width of these to one place.

-spec signature_type(signature_type()) -> integer().
signature_type(certificate_timestamp) -> 0;
signature_type(tree_hash) -> 1;
signature_type(test) -> 2.

-spec entry_type(entry_type()) -> integer().
entry_type(x509) -> 0;
entry_type(precert) -> 1;
entry_type(test) -> 2.

-spec leaf_type(leaf_type()) -> integer().
leaf_type(timestamped_entry) -> 0;
leaf_type(test) -> 1.    

-spec hash_alg_type(hash_alg_type()) -> integer().
hash_alg_type(none) -> 0;
hash_alg_type(md5) -> 1;
hash_alg_type(sha1) -> 2;
hash_alg_type(sha224) -> 3;
hash_alg_type(sha256) -> 4;
hash_alg_type(sha384) -> 5;
hash_alg_type(sha512) -> 6.

-spec signature_alg_type(signature_alg_type()) -> integer().
signature_alg_type(anonymous) -> 0;
signature_alg_type(rsa) -> 1;
signature_alg_type(dsa) -> 2;
signature_alg_type(ecdsa) -> 3.

-spec timestamp(now | integer()) -> integer().
timestamp(Timestamp) ->
    case Timestamp of
        now ->
            {NowMegaSec, NowSec, NowMicroSec} = now(),
            trunc(NowMegaSec * 1.0e9
                  + NowSec * 1.0e3
                  + NowMicroSec / 1.0e3);
        _ -> Timestamp
    end.

-spec serialise(plop_entry() | timestamped_entry() | mtl() |
                spt() | spt_signed() | sth() | sth_signed() |
                sig_and_hash_alg() | signature()) -> binary().
serialise(#plop_entry{
             type = TypeAtom,
             data = Data
             }) ->
    EntryType = entry_type(TypeAtom),
    <<EntryType:16, Data/binary>>;
serialise(#timestamped_entry{
             timestamp = Timestamp,
             entry = PlopEntry
            }) ->
    list_to_binary([<<Timestamp:64>>, serialise(PlopEntry)]);
serialise(#spt{
             version = Version,
             logid = LogID, 
             timestamp = Timestamp,
             signature = Signature
            }) ->
    list_to_binary([<<Version:8, LogID/binary, Timestamp:64>>,
                    serialise(Signature)]);
serialise(#spt_signed{
             version = Version,
             signature_type = SigtypeAtom,
             timestamp = Timestamp,
             entry_type = EntrytypeAtom,
             signed_entry = Entry
            }) ->
    Sigtype = signature_type(SigtypeAtom),
    Entrytype = entry_type(EntrytypeAtom),
    <<Version:8,
      Sigtype:8,
      Timestamp:64,
      Entrytype:16,
      Entry/binary>>;
serialise(#mtl{                                 % Merkle Tree Leaf.
             version = Version,
             leaf_type = TypeAtom,
             entry = TimestampedEntry
            }) ->
    LeafType = leaf_type(TypeAtom),
    list_to_binary([<<Version:8, LeafType:8>>, serialise(TimestampedEntry)]);
serialise(#sth_signed{                          % Signed Tree Head.
             version = Version,
             signature_type = SigtypeAtom,
             timestamp = Timestamp,
             tree_size = Treesize,
             root_hash = Roothash
            }) ->
    Sigtype = signature_type(SigtypeAtom),
    <<Version:8,
      Sigtype:8,
      Timestamp:64,
      Treesize:64,
      Roothash/binary>>;
serialise(#sig_and_hash_alg{
             hash_alg = HashAlgType,
             signature_alg = SignatureAlgType
            }) ->
    HashAlg = hash_alg_type(HashAlgType),
    SignatureAlg = signature_alg_type(SignatureAlgType),
    <<HashAlg:8, SignatureAlg:8>>;
serialise(#signature{
             algorithm = Algorithm,
             signature = Signature              % DER encoded.
             }) ->
    %% Encode a DSS signature according to RFC5246 section 4.7 and
    %% don't forget that the signature is a vector as specified in
    %% section 4.3 and has a length field.
    SigLen = size(Signature),
    list_to_binary([serialise(Algorithm),
                    <<SigLen:16, Signature/binary>>]).

%%%%%%%%%%%%%%%%%%%%
%% Internal tests. For more tests see ../test/.
serialise_test_() ->
    [?_assertEqual(
        <<0:8, 0:8, 0:64, 0:16, "foo">>,
        serialise(#spt_signed{
                     version = 0,
                     signature_type = certificate_timestamp,
                     timestamp = 0,
                     entry_type = x509,
                     signed_entry = <<"foo">>}))].
add_test() ->
    {ok, S} = init([?TESTPRIVKEYFILE, ?TESTPUBKEYFILE]),

    Data1 = <<"some data">>,
    {_Tree, SPT} =
        do_add(#timestamped_entry{
                  timestamp = 4711,
                  entry = #plop_entry{type = test, data = Data1}},
               S#state.privkey,
               S#state.logid),
    {_Tree1, SPT1} =
        do_add(#timestamped_entry{
                  timestamp = 4712,
                  entry = #plop_entry{type = test, data = Data1}},
               S#state.privkey,
               S#state.logid),
    ?assertEqual(SPT, SPT1),

    TE = #timestamped_entry{
            timestamp = 0,
            entry = #plop_entry{type = test, data = <<"some data">>}},
    SPTeq1 = spt(S#state.logid, S#state.privkey, TE),
    SPTeq2 = spt(S#state.logid, S#state.privkey, TE),
    ?assertNotEqual(SPTeq1, SPTeq2),          % DSA signatures differ!
    ok.

%% add_random_data(N, Size) ->
%%     lists:foreach(
%%       fun(_) -> plop:add(#timestamped_entry
%%                          {entry = #plop_entry
%%                           {type = test,
%%                            data = crypto:rand_bytes(Size)}}) end,
%%       lists:seq(1, N)).