%%% @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, the data and a hash
%%% of it is stored in a way that [mumble FIXME and FIXME]. In return
%%% you will get a proof of your entry being included in the log. This
%%% proof can later, together with the public key of the log, be used
%%% to prove 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([sth/1]).
-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(PLOPVERSION, 1).
-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()}).

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

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

%%%%%%%%%%%%%%%%%%%%
init([PubKeyfile, PrivKeyfile]) ->
    %% 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(PubKeyfile, PrivKeyfile),
    LogID = crypto:hash(sha256, public_key:der_encode(
                                  'ECPoint',
                                  element(2, element(1, Public_key)))), % FIXME!
    {ok, #state{pubkey = Public_key,
                privkey = Private_key,
                logid = LogID,
                hashtree = ht:create()}}.

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, []}).
sth(Data) ->
    gen_server:call(?MODULE, {sth, Data}).

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 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, data=~p~n",
                      [I, E#timestamped_entry.entry#plop_entry.data]),
            Record = Record#plop{           % DB consistency checking.
                       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,data=~p~n",
                      [ht:size(Tree) + 1, PlopEntry#plop_entry.data]),
            %% io:format("Creating new entry: index=~p~ndb hash: ~p~nMTL: ~p~nSPT: ~p~n",
            %%           [ht:size(Tree) + 1, DB_hash, MTL, NewSPT]),
            DB_data = #plop{index = ht:size(Tree) + 1,
                            hash = DB_hash,
                            mtl = MTL,
                            spt = NewSPT},
            db:add(DB_data),
            {ht:append(Tree, serialise(MTL)),   % New tree.
             NewSPT}                            % New SPT.
    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{
                                    version = 1,
                                    signature_type = certificate_timestamp,
                                    timestamp = Timestamp,
                                    entry_type = EntryType,
                                    signed_entry = EntryData})),
    Signature = signhash(BinToSign, PrivKey),
    #spt{
       version = ?PLOPVERSION,
       logid = LogID,
       timestamp = Timestamp,
       signature = Signature}.

%% @doc Signed Tree Head as specified in RFC6962 section 3.2.
sth(PrivKey, Tree, []) ->
    sth(PrivKey, Tree, #sth{timestamp = now});
sth(PrivKey, Tree, #sth{version = Version, timestamp = Timestamp_in}) ->
    Timestamp = timestamp(Timestamp_in),
    Treesize = ht:size(Tree),
    Roothash = ht:tree_hash(Tree),
    BinToSign =
        list_to_binary(serialise(#sth{version = Version,
                                      signature_type = tree_hash,
                                      timestamp = Timestamp,
                                      tree_size = Treesize,
                                      root_hash = Roothash})),
    Signature = signhash(BinToSign, PrivKey),
    STH = <<Treesize:64,
            Timestamp:64,
            Roothash/binary,
            Signature/binary>>,
    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 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() | spt() | spt_signed() | mtl() | sth()) -> 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,
       Signature/binary>>];
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 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>>].


%%%%%%%%%%%%%%%%%%%%
%% Tests.
serialise_test_() ->
    [?_assertEqual(
        <<1:8, 0:8, 0:64, 0:16, "foo">>,
        list_to_binary(serialise(#spt_signed{
                                    version = 1,
                                    signature_type = certificate_timestamp,
                                    timestamp = 0,
                                    entry_type = x509,
                                    signed_entry = <<"foo">>})))].
%%add_test_() ->
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!