%%% @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.

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

%% API.
-export([start_link/2, stop/0]).
-export([get_logid/0, serialise/1]).
-export([add/1, sth/0]).
%% 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("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(),
                hashtree :: ht:head()}).

%% @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 :: pos_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 :: pos_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]) ->
    io:format("plop starting~n"),
    %% 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!
    %% WARNING FIXME: The building of the tree is immensely expensive
    %% and slow -- don't do this with more than a couple of hundred of
    %% entries!
    {ok, #state{pubkey = Public_key,
                privkey = Private_key,
                logid = LogID,
                hashtree = build_tree_from_db()}}.

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_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,
                                  hashtree = Tree}) ->
    TimestampedEntry = #timestamped_entry{
                          timestamp = timestamp(Timestamp_in),
                          entry = Entry},
    {NewTree, SPT} = do_add(TimestampedEntry, Privkey, LogID, Tree),
    {reply, SPT, State#state{hashtree = NewTree}};

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

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

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

%%%%%%%%%%%%%%%%%%%%
-spec build_tree_from_db() -> ht:head().
build_tree_from_db() ->
    build_tree(ht:create(), lists:seq(0, db:size() - 1)).
-spec build_tree(ht:head(), list()) -> ht:head().
build_tree(Tree, []) ->
    Tree;
build_tree(Tree, [H|T]) ->
    Data = db_get_single_entry(H),
    build_tree(ht:append(Tree, Data), T).

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(), any()) -> {any(), binary()}.
do_add(TimestampedEntry = #timestamped_entry{entry = PlopEntry},
       Privkey, LogID, Tree) ->
    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}}},
            {Tree, SPT};              % State not changed, cached SPT.
        [] ->
            NewSPT = spt(LogID, Privkey, TimestampedEntry),
            MTL = #mtl{entry = TimestampedEntry},
            io:format("Creating new entry: index=~p~n", [ht:size(Tree)]),
            DB_data = #plop{index = ht:size(Tree),
                            hash = DB_hash,
                            mtl = MTL,
                            spt = NewSPT},
            {atomic, ok} = db:add(DB_data),
            {ht:append(Tree, serialise(MTL)),   % New tree.
             NewSPT};                           % New SPT.
        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 =
        list_to_binary(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'{}, ht:head(), sth_signed() | list()) -> sth().
sth(PrivKey, Tree, []) ->
    sth(PrivKey, Tree, #sth_signed{timestamp = now});
sth(PrivKey, Tree, #sth_signed{version = Version, timestamp = Timestamp_in}) ->
    Timestamp = timestamp(Timestamp_in),
    Treesize = ht:size(Tree),
    Roothash = ht:tree_hash(Tree),
    BinToSign =
        list_to_binary(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.

%% 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)}.

%% 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()) -> iolist().
serialise(#plop_entry{
             type = TypeAtom,
             data = Data
             }) ->
    EntryType = entry_type(TypeAtom),
    [<<EntryType:16,
       Data/binary>>];
serialise(#timestamped_entry{
             timestamp = Timestamp,
             entry = PlopEntry
            }) ->
    [<<Timestamp:64>>,
     serialise(PlopEntry)];
serialise(#spt{
             version = Version,
             logid = LogID, 
             timestamp = Timestamp,
             signature = Signature
            }) ->
    [<<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),
    [<<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),
    [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">>,
        list_to_binary(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,
               S#state.hashtree),
    {_Tree1, SPT1} =
        do_add(#timestamped_entry{
                  timestamp = 4712,
                  entry = #plop_entry{type = test, data = Data1}},
               S#state.privkey,
               S#state.logid,
               S#state.hashtree),
    ?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!

add_more_test() ->
    D = crypto:rand_bytes(3072),
    lists:foreach(
      fun(D) -> plop:add(#timestamped_entry
                         {entry = #plop_entry{type = test, data = D}}) end,
      lists:seq(0, 200)).