sha256 Hashing Enabled in ODBC Tables - Fixing MOD_REST and sha256 at the same time

-module(cyrsasl).
-author('alexey@process-one.net').

-export([start/0,
register_mechanism/3,
listmech/1,
server_new/7,
server_start/3,
server_step/2]).

-include("ejabberd.hrl").

-record(sasl_mechanism, {mechanism, module, password_type}).
-record(sasl_state, {service, myname, realm,
     get_password, check_password, check_password_digest,
     mech_mod, mech_state}).

-export([behaviour_info/1]).

behaviour_info(callbacks) ->
    [{mech_new, 4}, {mech_step, 2}];
behaviour_info(_Other) ->
    undefined.

start() ->
    ets:new(sasl_mechanism, [named_table,
     public,
     {keypos, #sasl_mechanism.mechanism}]),
    cyrsasl_plain:start([]),
%%    cyrsasl_digest:start([]),
%%    cyrsasl_scram:start([]),
%%    cyrsasl_anonymous:start([]),
    ok.

register_mechanism(Mechanism, Module, PasswordType) ->
    ets:insert(sasl_mechanism,
       #sasl_mechanism{mechanism = Mechanism,
       module = Module,
       password_type = PasswordType}).

%%% TODO: use callbacks
%%-include("ejabberd.hrl").
%%-include("jlib.hrl").
%%check_authzid(_State, Props) ->
%%    AuthzId = xml:get_attr_s(authzid, Props),
%%    case jlib:string_to_jid(AuthzId) of
%% error ->
%%     {error, "invalid-authzid"};
%% JID ->
%%     LUser = jlib:nodeprep(xml:get_attr_s(username, Props)),
%%     {U, S, R} = jlib:jid_tolower(JID),
%%     case R of
%% "" ->
%%     {error, "invalid-authzid"};
%% _ ->
%%     case {LUser, ?MYNAME} of
%% {U, S} ->
%%     ok;
%% _ ->
%%     {error, "invalid-authzid"}
%%     end
%%     end
%%    end.

check_credentials(_State, Props) ->
    User = xml:get_attr_s(username, Props),
    case jlib:nodeprep(User) of
error ->
    {error, "not-authorized"};
"" ->
    {error, "not-authorized"};
_LUser ->
    ok
    end.

listmech(Host) ->
    Mechs = ets:select(sasl_mechanism,
       [{#sasl_mechanism{mechanism = '$1',
password_type = '$2',
_ = '_'},
case catch ejabberd_auth:store_type(Host) of
external ->
      [{'==', '$2', plain}];
scram ->
      [{'/=', '$2', digest}];
{'EXIT',{undef,[{Module,store_type,[]} | _]}} ->
      ?WARNING_MSG("~p doesn't implement the function store_type/0", [Module]),
      [];
_Else ->
      []
end,
['$1']}]),
    filter_anonymous(Host, Mechs).

server_new(Service, ServerFQDN, UserRealm, _SecFlags,
   GetPassword, CheckPassword, CheckPasswordDigest) ->
    #sasl_state{service = Service,
myname = ServerFQDN,
realm = UserRealm,
get_password = GetPassword,
check_password = CheckPassword,
check_password_digest= CheckPasswordDigest}.

server_start(State, Mech, ClientIn) ->
    case lists:member(Mech, listmech(State#sasl_state.myname)) of
true ->
    case ets:lookup(sasl_mechanism, Mech) of
[#sasl_mechanism{module = Module}] ->
    {ok, MechState} = Module:mech_new(
State#sasl_state.myname,
State#sasl_state.get_password,
State#sasl_state.check_password,
State#sasl_state.check_password_digest),
    server_step(State#sasl_state{mech_mod = Module,
mech_state = MechState},
ClientIn);
_ ->
    {error, "no-mechanism"}
    end;
false ->
    {error, "no-mechanism"}
    end.

server_step(State, ClientIn) ->
    Module = State#sasl_state.mech_mod,
    MechState = State#sasl_state.mech_state,
    case Module:mech_step(MechState, ClientIn) of
{ok, Props} ->
    case check_credentials(State, Props) of
ok ->
    {ok, Props};
{error, Error} ->
    {error, Error}
    end;
{ok, Props, ServerOut} ->
    case check_credentials(State, Props) of
ok ->
    {ok, Props, ServerOut};
{error, Error} ->
    {error, Error}
    end;
{continue, ServerOut, NewMechState} ->
    {continue, ServerOut,
     State#sasl_state{mech_state = NewMechState}};
{error, Error, Username} ->
    {error, Error, Username};
{error, Error} ->
    {error, Error}
    end.

%% Remove the anonymous mechanism from the list if not enabled for the given
%% host
filter_anonymous(Host, Mechs) ->
    case ejabberd_auth_anonymous:is_sasl_anonymous_enabled(Host) of
true  -> Mechs;
false -> Mechs -- ["ANONYMOUS"]
    end.

-module(ejabberd_auth_odbc).
-author('alexey@process-one.net').

%% External exports
-export([start/1,
set_password/3,
check_password/3,
check_password/5,
try_register/3,
dirty_get_registered_users/0,
get_vh_registered_users/1,
get_vh_registered_users/2,
get_vh_registered_users_number/1,
get_vh_registered_users_number/2,
get_password/2,
get_password_s/2,
is_user_exists/2,
remove_user/2,
remove_user/3,
store_type/0,
plain_password_required/0
]).

-include("ejabberd.hrl").

%%%----------------------------------------------------------------------
%%% API
%%%----------------------------------------------------------------------
start(_Host) ->
    ok.

hash(User, Password) ->
    Split="--Zach Calvert--",
    bitstring_to_list(base64:encode(erlsha2:sha256((string:concat(string:concat(Password, Split), User))))).

plain_password_required() ->
    false.

store_type() ->
plain.

%% @spec (User, Server, Password) -> true | false | {error, Error}
check_password(User, Server, Password) ->
    HashedPassword = hash(User, Password),
    case jlib:nodeprep(User) of
error ->
    false;
LUser ->
    Username = ejabberd_odbc:escape(LUser),
    LServer = jlib:nameprep(Server),
    try odbc_queries:get_password(LServer, Username) of
{selected, ["password"], [{HashedPassword}]} ->
    Password /= ""; %% Password is correct, and not empty
{selected, ["password"], [{_Password2}]} ->
    false; %% Password is not correct
{selected, ["password"], []} ->
    false; %% Account does not exist
{error, _Error} ->
    false %% Typical error is that table doesn't exist
    catch
_:_ ->
    false %% Typical error is database not accessible
    end
    end.

%% @spec (User, Server, Password, Digest, DigestGen) -> true | false | {error, Error}
check_password(User, Server, Password, Digest, DigestGen) ->
    HashedPassword = hash(User, Password),
    case jlib:nodeprep(User) of
error ->
    false;
LUser ->
    Username = ejabberd_odbc:escape(LUser),
    LServer = jlib:nameprep(Server),
    try odbc_queries:get_password(LServer, Username) of
%% Account exists, check if password is valid
{selected, ["password"], [{HashedPassword}]} ->
    Password /= ""; %% Password is correct, and not empty
{selected, ["password"], [{Passwd}]} ->
    DigRes = if
Digest /= "" ->
     Digest == DigestGen(Passwd);
true ->
     false
     end,
    if DigRes ->
    true;
       true ->
    (Passwd == HashedPassword) and (Password /= "")
    end;
{selected, ["password"], []} ->
    false; %% Account does not exist
{error, _Error} ->
    false %% Typical error is that table doesn't exist
    catch
_:_ ->
    false %% Typical error is database not accessible
    end
    end.

%% @spec (User::string(), Server::string(), Password::string()) ->
%%       ok | {error, invalid_jid}
set_password(User, Server, Password) ->
    case jlib:nodeprep(User) of
error ->
    {error, invalid_jid};
LUser ->
    Username = ejabberd_odbc:escape(LUser),
            HashedPassword = hash(User, Password),
    LServer = jlib:nameprep(Server),
    case catch odbc_queries:set_password_t(LServer, Username, HashedPassword) of
        {atomic, ok} -> ok;
        Other -> {error, Other}
    end
    end.

%% @spec (User, Server, Password) -> {atomic, ok} | {atomic, exists} | {error, invalid_jid}
try_register(User, Server, Password) ->
    case jlib:nodeprep(User) of
error ->
    {error, invalid_jid};
LUser ->
    Username = ejabberd_odbc:escape(LUser),
            HashedPassword = hash(User, Password),
    LServer = jlib:nameprep(Server),
    case catch odbc_queries:add_user(LServer, Username, HashedPassword) of
{updated, 1} ->
    {atomic, ok};
_ ->
    {atomic, exists}
    end
    end.

dirty_get_registered_users() ->
    Servers = ejabberd_config:get_vh_by_auth_method(odbc),
    lists:flatmap(
      fun(Server) ->
      get_vh_registered_users(Server)
      end, Servers).

get_vh_registered_users(Server) ->
    LServer = jlib:nameprep(Server),
    case catch odbc_queries:list_users(LServer) of
{selected, ["username"], Res} ->
    [{U, LServer} || {U} <- Res];
_ ->
    []
    end.

get_vh_registered_users(Server, Opts) ->
    LServer = jlib:nameprep(Server),
    case catch odbc_queries:list_users(LServer, Opts) of
{selected, ["username"], Res} ->
    [{U, LServer} || {U} <- Res];
_ ->
    []
    end.

get_vh_registered_users_number(Server) ->
    LServer = jlib:nameprep(Server),
    case catch odbc_queries:users_number(LServer) of
{selected, [_], [{Res}]} ->
    list_to_integer(Res);
_ ->
    0
    end.

get_vh_registered_users_number(Server, Opts) ->
    LServer = jlib:nameprep(Server),
    case catch odbc_queries:users_number(LServer, Opts) of
{selected, [_], [{Res}]} ->
    list_to_integer(Res);
_Other ->
    0
    end.

get_password(User, Server) ->
    case jlib:nodeprep(User) of
error ->
    false;
LUser ->
    Username = ejabberd_odbc:escape(LUser),
    LServer = jlib:nameprep(Server),
    case catch odbc_queries:get_password(LServer, Username) of
{selected, ["password"], [{Password}]} ->
    Password;
_ ->
    false
    end
    end.

get_password_s(User, Server) ->
    case jlib:nodeprep(User) of
error ->
    "";
LUser ->
    Username = ejabberd_odbc:escape(LUser),
    LServer = jlib:nameprep(Server),
    case catch odbc_queries:get_password(LServer, Username) of
{selected, ["password"], [{Password}]} ->
    Password;
_ ->
    ""
    end
    end.

%% @spec (User, Server) -> true | false | {error, Error}
is_user_exists(User, Server) ->
    case jlib:nodeprep(User) of
error ->
    false;
LUser ->
    Username = ejabberd_odbc:escape(LUser),
    LServer = jlib:nameprep(Server),
    try odbc_queries:get_password(LServer, Username) of
{selected, ["password"], [{_Password}]} ->
    true; %% Account exists
{selected, ["password"], []} ->
    false; %% Account does not exist
{error, Error} ->
    {error, Error} %% Typical error is that table doesn't exist
    catch
_:B ->
    {error, B} %% Typical error is database not accessible
    end
    end.

%% @spec (User, Server) -> ok | error
%% @doc Remove user.
%% Note: it may return ok even if there was some problem removing the user.
remove_user(User, Server) ->
    case jlib:nodeprep(User) of
error ->
    error;
LUser ->
    Username = ejabberd_odbc:escape(LUser),
    LServer = jlib:nameprep(Server),
    catch odbc_queries:del_user(LServer, Username),
ok
    end.

%% @spec (User, Server, Password) -> ok | error | not_exists | not_allowed
%% @doc Remove user if the provided password is correct.
remove_user(User, Server, Password) ->
    case jlib:nodeprep(User) of
error ->
    error;
LUser ->
    Username = ejabberd_odbc:escape(LUser),
            HashedPassword = hash(User, Password),
    LServer = jlib:nameprep(Server),
    F = fun() ->
Result = odbc_queries:del_user_return_password(
   LServer, Username, HashedPassword),
case Result of
    {selected, ["password"], [{Password}]} ->
ok;
    {selected, ["password"], []} ->
not_exists;
    _ ->
not_allowed
end
end,
    {atomic, Result} = odbc_queries:sql_transaction(LServer, F),
    Result
    end.