#!/usr/bin/env escript
%%! -args_file /var/riak/ssl_distribution.args_file
%% -*- mode: erlang;erlang-indent-level: 4;indent-tabs-mode: nil -*-
%% ex: ft=erlang ts=4 sw=4 et
%% -------------------------------------------------------------------
%%
%% nodetool: Helper Script for interacting with live nodes
%%
%% -------------------------------------------------------------------
-mode(compile).

main(Args) ->
    ok = start_epmd(),
    %% Extract the args
    {RestArgs, TargetNode} = process_args(Args, [], undefined),

    %% process_args() has side-effects (e.g. when processing "-name"),
    %% so take care of app-starting business first.
    [application:start(App) || App <- [crypto, public_key, ssl]],

    %% any commands that don't need a running node
    case RestArgs of
        ["chkconfig", File] ->
            case file:consult(File) of
                {ok, Terms} ->
                    case validate(Terms) of
                        ok ->
                            io:format("ok\n"),
                            halt(0);
                        {error, Problems} ->
                            lists:foreach(fun print_issue/1, Problems),
                            %% halt(1) if any problems were errors
                            halt(case [x || {error, _} <- Problems] of
                                     [] -> 0;
                                     _  -> 1
                                 end)
                    end;
                {error, {Line, Mod, Term}} ->
                    io:format(standard_error, ["Error on line ", file:format_error({Line, Mod, Term}), "\n"], []),
                    halt(1);
                {error, R} ->
                    io:format(standard_error, ["Error reading config file: ", file:format_error(R), "\n"], []),
                    halt(1)
            end;
        _ ->
            ok
    end,

    %% See if the node is currently running  -- if it's not, we'll bail
    case {net_kernel:hidden_connect_node(TargetNode), net_adm:ping(TargetNode)} of
        {true, pong} ->
            ok;
        {false,pong} ->
            io:format("Failed to connect to node ~p .\n", [TargetNode]),
            halt(1);
        {_, pang} ->
            io:format("Node ~p not responding to pings.\n", [TargetNode]),
            halt(1)
    end,

    case RestArgs of
        ["getpid"] ->
            io:format("~p\n", [list_to_integer(rpc:call(TargetNode, os, getpid, []))]);
        ["ping"] ->
            %% If we got this far, the node already responsed to a ping, so just dump
            %% a "pong"
            io:format("pong\n");
        ["stop"] ->
            io:format("~p\n", [rpc:call(TargetNode, init, stop, [], 60000)]);
        ["restart"] ->
            io:format("~p\n", [rpc:call(TargetNode, init, restart, [], 60000)]);
        ["reboot"] ->
            io:format("~p\n", [rpc:call(TargetNode, init, reboot, [], 60000)]);
        ["rpc", Module, Function | RpcArgs] ->
            case rpc:call(TargetNode, list_to_atom(Module), list_to_atom(Function),
                          [RpcArgs], 60000) of
                ok ->
                    ok;
                {badrpc, Reason} ->
                    io:format("RPC to ~p failed: ~p\n", [TargetNode, Reason]),
                    halt(1);
                _ ->
                    halt(1)
            end;
        ["rpc_infinity", Module, Function | RpcArgs] ->
            case rpc:call(TargetNode, list_to_atom(Module), list_to_atom(Function), [RpcArgs], infinity) of
                ok ->
                    ok;
                {badrpc, Reason} ->
                    io:format("RPC to ~p failed: ~p\n", [TargetNode, Reason]),
                    halt(1);
                _ ->
                    halt(1)
            end;
        ["rpcterms", Module, Function, ArgsAsString] ->
            case rpc:call(TargetNode, list_to_atom(Module), list_to_atom(Function),
                          consult(ArgsAsString), 60000) of
                {badrpc, Reason} ->
                    io:format("RPC to ~p failed: ~p\n", [TargetNode, Reason]),
                    halt(1);
                Other ->
                    io:format("~p\n", [Other])
            end;
        Other ->
            io:format("Other: ~p\n", [Other]),
            io:format("Usage: nodetool {chkconfig|getpid|ping|stop|restart|reboot|rpc|rpc_infinity|rpcterms}\n")
    end,
    net_kernel:stop().


unique([]) ->
    true;
unique([H|T]) ->
    case lists:all(fun(X) -> X =/= H end, T) of
        true -> unique(T);
        false -> false
    end.

process_args([], Acc, TargetNode) ->
    {lists:reverse(Acc), TargetNode};
process_args(["-setcookie", Cookie | Rest], Acc, TargetNode) ->
    erlang:set_cookie(node(), list_to_atom(Cookie)),
    process_args(Rest, Acc, TargetNode);
process_args(["-name", TargetName | Rest], Acc, _) ->
    ThisNode = append_node_suffix(TargetName, "_maint_"),
    {ok, _} = net_kernel:start([ThisNode, longnames]),
    process_args(Rest, Acc, nodename(TargetName));
process_args(["-sname", TargetName | Rest], Acc, _) ->
    ThisNode = append_node_suffix(TargetName, "_maint_"),
    {ok, _} = net_kernel:start([ThisNode, shortnames]),
    process_args(Rest, Acc, nodename(TargetName));
process_args([Arg | Rest], Acc, Opts) ->
    process_args(Rest, [Arg | Acc], Opts).


start_epmd() ->
    [] = os:cmd(epmd_path() ++ " -daemon"),
    ok.

epmd_path() ->
    ErtsBinDir = filename:dirname(escript:script_name()),
    Name = "epmd",
    case os:find_executable(Name, ErtsBinDir) of
        false ->
            case os:find_executable(Name) of
                false ->
                    io:format("Could not find epmd.~n"),
                    halt(1);
                GlobalEpmd ->
                    GlobalEpmd
            end;
        Epmd ->
            Epmd
    end.


nodename(Name) ->
    case string:tokens(Name, "@") of
        [_Node, _Host] ->
            list_to_atom(Name);
        [Node] ->
            [_, Host] = string:tokens(atom_to_list(node()), "@"),
            list_to_atom(lists:concat([Node, "@", Host]))
    end.

append_node_suffix(Name, Suffix) ->
    case string:tokens(Name, "@") of
        [Node, Host] ->
            list_to_atom(lists:concat([Node, Suffix, os:getpid(), "@", Host]));
        [Node] ->
            list_to_atom(lists:concat([Node, Suffix, os:getpid()]))
    end.


%%
%% Given a string or binary, parse it into a list of terms, ala file:consult/0
%%
consult(Str) when is_list(Str) ->
    consult([], Str, []);
consult(Bin) when is_binary(Bin)->
    consult([], binary_to_list(Bin), []).

consult(Cont, Str, Acc) ->
    case erl_scan:tokens(Cont, Str, 0) of
        {done, Result, Remaining} ->
            case Result of
                {ok, Tokens, _} ->
                    {ok, Term} = erl_parse:parse_term(Tokens),
                    consult([], Remaining, [Term | Acc]);
                {eof, _Other} ->
                    lists:reverse(Acc);
                {error, Info, _} ->
                    {error, Info}
            end;
        {more, Cont1} ->
            consult(Cont1, eof, Acc)
    end.


%%
%% Validation functions for checking the app.config
%%
validate([Terms]) ->

    Results = [ValidateFun(Terms) || ValidateFun <- get_validation_funs()],
    Failures = [Res || Res <- Results, Res /= true],
    case Failures of
        [] ->
            ok;
        _ ->
            {error, Failures}
    end.

%% Some initial and basic checks for the app.config file
get_validation_funs() ->
    [fun handoffport_exists/1,
     fun pbc_port_exists/1,
     fun have_unique_ports/1].


handoffport_exists(Terms) ->
    case get_env(Terms, riak_core, handoff_port) of
        undefined ->
            {error, "required handoff_port is not defined"};
        _ ->
            true
    end.


pbc_port_exists(Terms) ->
    case get_env(Terms, riak_api, pb_port) of
        undefined ->
            {warning, "pb_port undefined, protobuffs interface not enabled"};
        _ ->
            true
    end.


%% Check if any set ports have the same value
have_unique_ports(Terms) ->
    HandoffPort = get_env(Terms, riak_core, handoff_port),
    PBCPort = get_env(Terms, riak_api, pb_port),
    HTTP = get_env(Terms, riak_core, http),
    HTTPS = get_env(Terms, riak_core, https),

    Ports = [get_setting_if_exists(PortSetting) ||
                PortSetting <- [HandoffPort, PBCPort, HTTP, HTTPS]],

    % Remove empty lists so unique doesn't match on a empty lists
    SetPorts = lists:filter(fun(P) -> P /= [] end, Ports),

    case unique(SetPorts) of
        true -> true;
        false -> {error, "Port settings are not all unique"}
    end.


get_setting_if_exists(undefined) -> [];
get_setting_if_exists({ok, [{_, Setting}]}) -> [Setting];
get_setting_if_exists({ok, Setting}) -> [Setting].


get_env(List, Section, Setting) ->
    case [{X,Y} || {X,Y} <- List, X =:= Section] of
        [{Section, Settings}] ->
            case [{X,Y} || {X,Y} <- Settings, X =:= Setting] of
                [{Setting, Val}] -> {ok, Val};
                _ -> undefined
            end;
        _ -> undefined
    end.


print_issue({warning, Warning}) ->
    io:format(standard_error, "Warning in app.config: ~s~n", [Warning]);
print_issue({error, Error}) ->
    io:format(standard_error, "Error in app.config: ~s~n", [Error]).
