gistfile1.erl

%%%-------------------------------------------------------------------
%%% File    : mod_message_log.erl
%%% Author  : Holger Weiss <holger@zedat.fu-berlin.de>
%%% Purpose : Log one line per message transmission
%%% Created : 27 May 2014 by Holger Weiss <holger@zedat.fu-berlin.de>
%%%-------------------------------------------------------------------

-module(mod_message_log).
-author('holger@zedat.fu-berlin.de').

-behaviour(gen_mod).

-export([start/2,
	 stop/1,
	 init/1,
	 perform_upgrade/1]).

-export([log_packet_send/3,
	 log_packet_receive/4,
	 log_packet_offline/3,
	 reopen_log/0]).

-include("ejabberd.hrl").
-include("jlib.hrl").


-define(PROCNAME, ?MODULE).
-define(DEFAULT_HOST, "127.0.0.1").
-define(DEFAULT_PORT, 6379).
-define(DEFAULT_DB, 0).

-record(config, {redis_host = ?DEFAULT_HOST, redis_port = ?DEFAULT_PORT, redis_db = ?DEFAULT_DB, rclient}).

start(Host, Opts) ->
    ejabberd_hooks:add(user_send_packet, Host, ?MODULE,
		       log_packet_send, 42),
    ejabberd_hooks:add(user_receive_packet, Host, ?MODULE,
		       log_packet_receive, 42),
    ejabberd_hooks:add(offline_message_hook, Host, ?MODULE,
		       log_packet_offline, 42),
    case whereis(?PROCNAME) of
      undefined ->
	      ejabberd_hooks:add(reopen_log_hook, ?MODULE, reopen_log, 42),
        Redis_host = gen_mod:get_opt(redis_host, Opts, fun(V) -> V end, ?DEFAULT_HOST),
        Redis_port = gen_mod:get_opt(redis_port, Opts, fun(V) -> V end, ?DEFAULT_PORT),
        Redis_db = gen_mod:get_opt(redis_db, Opts, fun(V) -> V end, ?DEFAULT_DB),
        
        register(?PROCNAME, spawn(?MODULE, init,[#config{redis_host = Redis_host, redis_port = Redis_port, redis_db = Redis_db}])),
        ok;
      _ ->
	      ok
    end.

stop(Host) ->
    ejabberd_hooks:delete(user_send_packet, Host, ?MODULE,
			  log_packet_send, 42),
    ejabberd_hooks:delete(user_receive_packet, Host, ?MODULE,
			  log_packet_receive, 42),
    ejabberd_hooks:delete(offline_message_hook, Host, ?MODULE,
			  log_packet_offline, 42),
    case whereis(?PROCNAME) of
      undefined ->
	      ok;
      _ ->
	      ejabberd_hooks:delete(reopen_log_hook, ?MODULE, reopen_log, 42),
        gen_mod:get_module_proc(Host, ?PROCNAME) ! stop,
	      ok
    end.

init(Config) ->
    {ok, C} = eredis:start_link(Config#config.redis_host, Config#config.redis_port, Config#config.redis_db),
    loop(Config#config{rclient = C}).

log_packet_send(From, To, Packet) ->
    log_packet(outgoing, From, To, Packet).

log_packet_receive(JID, From, _To, Packet) ->
    log_packet(incoming, From, JID, Packet).

log_packet_offline(From, To, Packet) ->
    log_packet(offline, From, To, Packet).

reopen_log() ->
    ?PROCNAME ! reopen.

%% Internal functions.

log_packet(Direction, From, To, #xmlel{name = <<"message">>} = Packet) ->
    case xml:get_subtag(Packet, <<"body">>) of
        #xmlel{children = Body} when length(Body) > 0 ->
            MessageBody = get_message_body(Packet),
            Id = get_message_id(Packet),
	          ?PROCNAME ! {message, Direction, From, To, MessageBody, Id};
    _ ->
       ok
    end;
log_packet(_Direction, _From, _To, _Packet) ->
    ok.

get_message_body(Packet) ->
     xml:get_path_s(Packet, [{elem, list_to_binary("body")}, cdata]).

get_message_id(#xmlel{attrs = Attrs}) ->
    case xml:get_attr_s(<<"id">>, Attrs) of
      <<"">> -> 
        <<"unknown">>;
      Id ->
       Id
    end.

loop(Config) ->
    receive
      {message, Direction, From, To, MessageBody, Id} ->
        write_log(Config#config.rclient, Direction, From, To, MessageBody, Id),
        loop(Config);
      reopen ->
        eredis:stop(Config#config.rclient),
        {ok, C} = eredis:start_link(Config#config.redis_host, Config#config.redis_port, Config#config.redis_db),
        loop(Config#config{rclient = C});
      upgrade ->
        mod_message_log:perform_upgrade(Config);
      stop ->
        exit(normal)
    end.

write_log(C, Direction, From, To, MessageBody, Id) ->
    Timestamp = get_timestamp(),
    Record = io_lib:format("{\"timestamp\" : ~B, \"direction\" : ~s, \"from\" : ~s, \"to\" : ~s, \"msg\" : ~s, \"msg_id\" : ~s} ~n",
			   [Timestamp, Direction,
			    jlib:jid_to_string(From),
			    jlib:jid_to_string(To),
          MessageBody, Id
         ]),
    eredis:q(C, ["RPUSH", "ejabberd-chat-log", Record]).

get_timestamp() ->
  {Mega,Sec,Micro} = erlang:now(),
  (Mega*1000000+Sec)*1000 + Micro div 1000.


perform_upgrade(Config) ->
    loop(Config).