nodetool

#!/usr/bin/env escript
%% -*- erlang -*-
%%! -hidden
-mode(compile).
%% -------------------------------------------------------------------
%% Ansible nodetool module
%% Copyright (c) 2016 Roberto Aloi, Claes Wikstrom
%% License: MIT
%% -------------------------------------------------------------------
%% -*- mode: erlang;erlang-indent-level: 4;indent-tabs-mode: nil -*-
%% ex: ft=erlang ts=4 sw=4 et
%% -------------------------------------------------------------------
%%
%% DOCUMENTATION
%%
%% module:            nodetool
%% short_description: A little helper to interact with Erlang nodes
%% description:       A little helper to interact with Erlang nodes
%% author:            Roberto Aloi
%% requirements:      Erlang
%% options:
%%   node:
%%     description: The remote Erlang node
%%     required:    true
%%   action:
%%     description: The action to be performed
%%     choices:     [getpid, ping, stop, restart, reboot, eval]
%%     required:    true
%%   nametype:
%%     description: Nametype to be used
%%     choices:     [longnames, shortnames]
%%     required:    false
%%     default:     shortnames
%%   cookie:
%%     description: Erlang Cookie to be used for the connection
%%     required:    false
%%   timeout:
%%     description: Timeout (in ms) for the actions
%%     required:    false
%%     default:     60000
%%
%% -------------------------------------------------------------------
-define(RPC_TIMEOUT, 60 * 1000). %% 1 minute

main([ConfigFile]) ->

    %% Extract the args
    Args = parse_args(ConfigFile),

    Action   = strip(get_arg("action",   Args)),
    Node     = strip(get_arg("node",     Args)),
    NameType = list_to_existing_atom(get_arg("nametype", Args, "shortnames")),

    ThisNode   = append_node_suffix(Node, "_maint_"),

    ok = start_distribution(ThisNode, NameType),

    case get_arg("cookie", Args) of
        undefined ->
            ok;
        Cookie    ->
            erlang:set_cookie(node(), list_to_atom(strip_dq(Cookie)))
    end,

    TargetNode = nodename(Node),

    case Action of
        "getpid" ->
            rpc(os, getpid, [], Args, fun erlang:list_to_integer/1);
        "ping" ->
            case net_adm:ping(TargetNode) of
                pong ->
                    success("pong");
                pang ->
                    fail("Node ~p not responding to pings.", [TargetNode])
            end;
        "stop" ->
            rpc(init, stop, [], Args);
        "restart" ->
            rpc(init, restart, [], Args);
        "reboot" ->
            rpc(init, reboot, [], Args);
        "eval" ->
            Command0 = get_arg("command", Args, ""),
            Command  = strip(Command0) ++ ".",
            case erl_scan:string(Command) of
                {ok, Tokens, _} ->
                    case erl_parse:parse_exprs(Tokens) of
                        {ok, Exprs} ->
                            Callback = fun({value, Value, _}) -> Value end,
                            rpc(erl_eval, exprs, [Exprs, []], Args, Callback);
                        {error, Reason} ->
                            Msg = format_parse_error(Reason),
                            fail(Msg)
                    end;
                {error, Reason, _} ->
                    Msg = format_parse_error(Reason),
                    fail(Msg)
            end;
        _Other ->
            fail("Wrong usage")
    end,
    net_kernel:stop().

parse_args(ConfigFile) ->
    case file:read_file(ConfigFile) of
        {error, Reason} ->
            fail("~p", [Reason]);
        {ok, Config} ->
            parse_args(re:split(Config, " "), [])
    end.

parse_args([], Acc) ->
    Acc; %% No need to reverse
parse_args([H|T], Acc) ->
    case binary:split(H, <<"=">>) of
        [Key, Value] ->
            parse_args(T, [{binary_to_list(Key), binary_to_list(Value)}|Acc]);
        _ ->
            parse_args(T, Acc)
    end.

start_distribution(Name, NameType) ->
    ok = start_epmd(),
    application:set_env(net_kernel, dist_auto_connect, false), %% Hidden node
    {ok, _} = net_kernel:start([Name, NameType]),
    ok.

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 ->
                    fail("Could not find epmd.");
                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] ->
            String = lists:concat([Node, Suffix, os:getpid(), "@", Host]),
            list_to_atom(string:strip(String, both, $'));
        [Node] ->
            String = lists:concat([Node, Suffix, os:getpid()]),
            list_to_atom(string:strip(String, both, $'))
    end.

success(String) ->
    success(String, []).

success(String, Args) ->
    success(String, Args, "").

success(String, Args, StdOut) ->
    Msg = lists:flatten(io_lib:format(String, Args)),
    io:format("{\"changed\": true, \"remote_output\": ~p, \"stdout\": ~p, \"rc\": 0}", [StdOut, Msg]),
    halt(0).

fail(String) ->
    fail(String, []).

fail(String, Args) ->
    fail(String, Args, "").

fail(String, Args, StdOut) ->
    Msg = lists:flatten(io_lib:format(String, Args)),
    io:format("{\"failed\" : true, \"remote_output\": ~p, \"stdout\": ~p, \"rc\": 1}", [StdOut, Msg]),
    halt(1).

format_parse_error({_Line, Mod, Err}) ->
    lists:flatten(Mod:format_error(Err)).

get_arg(Key, Args) ->
    get_arg(Key, Args, undefined).

get_arg(Key, Args, Default) ->
    proplists:get_value(Key, Args, Default).

rpc(Mod, Fun, Args, Opts) ->
    rpc(Mod, Fun, Args, Opts, fun(X) -> X end).

rpc(Mod, Fun, Args, Opts, Callback) ->
    Node    = nodename(strip(get_arg("node", Opts))),
    Timeout = case get_arg("timeout", Opts) of
                undefined -> ?RPC_TIMEOUT;
                V         ->
                  Stripped = strip(V),
                  {N, _} = string:to_integer(Stripped),
                  N
              end,
    case rpc_capture_output(Node, Mod, Fun, Args, Timeout) of
        {{badrpc, Reason}, IO} ->
            fail("RPC to ~p failed: ~p", [Node, Reason], IO);
        {Result, IO} ->
            success("~p", [Callback(Result)], IO)
    end.

rpc_capture_output(Node, Mod, Fun, Args, Timeout) ->
    P  = spawn_link(fun simple_group_leader/0),
    GL = group_leader(),
    group_leader(P, self()),
    Result = rpc:call(Node, Mod, Fun, Args, Timeout),
    group_leader(GL, self()),
    P ! {flush_request, self()},
    receive
        {flush_response, IO} -> IO
    end,
    {Result, IO}.

strip(A) ->
  strip_sq(strip_dq(A)).

strip_sq(A) ->
  string:strip(A, both, $').

strip_dq(A) ->
  string:strip(A, both, $").

simple_group_leader() ->
    simple_group_leader([]).

simple_group_leader(Acc) ->
    receive
        {flush_request, From} ->
            From ! {flush_response, lists:flatten(lists:reverse(Acc))};
        {io_request, From, ReplyAs, {put_chars, Encoding, Characters}} ->
            List = unicode:characters_to_list(Characters, Encoding),
            From ! {io_reply, ReplyAs, ok},
            simple_group_leader([List|Acc]);
        {io_request, From, ReplyAs, {put_chars, Encoding, M, F, A}} ->
            String = erlang:apply(M, F, A),
            List = unicode:characters_to_list(String, Encoding),
            From ! {io_reply, ReplyAs, ok},
            simple_group_leader([List|Acc]);
        {io_request, From, ReplyAs, Request} ->
            From ! {io_reply, ReplyAs, {error, request}},
            fail("Unimplemented request type: ~p", [Request]);
        Other ->
            fail("Unexpected I/O message: ~p", [Other])
    end.