connect2beam: a script that finds Erlang nodes and connect to them

connect2beam

#!/usr/bin/python

# connect2beam is a script that finds Erlang nodes and connects to them or
# kills them.
#
# Example for connection:
#
#     $ connect2beam
#     Available nodes: index name (cookie, pid)
#
#         1. refactorerl@localhost (None 16549)
#         2. [email protected] (wombat 17692)
#         3. [email protected] (riak 54911)
#
#     [...]
#
#     > 3
#     ['/Users/hcs/w/riak/clusters/riak-ee-2.0.5/1/bin/../erts-5.10.3/bin/erl',
#      '-name', '[email protected]', '-remsh', '[email protected]',
#      '-setcookie', 'riak']
#     Erlang R16B02_basho6 (erts-5.10.3) [source-bcd8abb] [64-bit] [smp:4:4]
#     [async-threads:10] [kernel-poll:false] [frame-pointer] [dtrace]
#
#     Eshell V5.10.3  (abort with ^G)
#     ([email protected])1>
#
# Example for killing:
#
#     $ connect2beam
#
#     Available nodes: index name (cookie, pid)
#
#         1. refactorerl@localhost (None 16549)
#         2. [email protected] (wombat 17692)
#         3. [email protected] (riak 54911)
#
#     [...]
#
#     > k 2 3
#     ['kill', '22529']
#     ['kill', '22566']
#
# You can also connect to an Elixir node using the iex shell:
#
#     Available nodes: index name (cookie, pid)
#
#         1. [email protected] (cookietest 71097)
#
#         [...]
#
#     > x 1
#     ['iex', '--name', '[email protected]', '--remsh',
#     '[email protected]', '--cookie', 'cookietest']
#     Erlang/OTP 18 [erts-7.2.1] [source] [64-bit] [smp:4:4] [async-threads:10]
#     [kernel-poll:false]
#
#     Interactive Elixir (1.2.1) - press Ctrl+C to exit (type h() ENTER for help)
#     iex([email protected])1>

import re
import sys
import subprocess
import os
import signal
import random


# Copied from http://stacyprowell.com/blog/2009/03/trapping-ctrlc-in-python/
class BreakHandler:
    '''
    Trap CTRL-C, set a flag, and keep going.  This is very useful for
    gracefully exiting database loops while simulating transactions.

    To use this, make an instance and then enable it.  You can check
    whether a break was trapped using the trapped property.

    # Create and enable a break handler.
    ih = BreakHandler()
    ih.enable()
    for x in big_set:
        complex_operation_1()
        complex_operation_2()
        complex_operation_3()
        # Check whether there was a break.
        if ih.trapped:
            # Stop the loop.
            break
    ih.disable()
    # Back to usual operation...
    '''

    def __init__(self, emphatic=9):
        '''
        Create a new break handler.

        @param emphatic: This is the number of times that the user must
                    press break to *disable* the handler.  If you press
                    break this number of times, the handler is automagically
                    disabled, and one more break will trigger an old
                    style keyboard interrupt.  The default is nine.  This
                    is a Good Idea, since if you happen to lose your
                    connection to the handler you can *still* disable it.
        '''
        self._count = 0
        self._enabled = False
        self._emphatic = emphatic
        self._oldhandler = None
        return

    def _reset(self):
        '''
        Reset the trapped status and count.  You should not need to use this
        directly; instead you can disable the handler and then re-enable it.
        This is better, in case someone presses CTRL-C during this operation.
        '''
        self._count = 0
        return

    def enable(self):
        '''
        Enable trapping of the break.  This action also resets the
        handler count and trapped properties.
        '''
        if not self._enabled:
            self._reset()
            self._enabled = True
            self._oldhandler = signal.signal(signal.SIGINT, self)
        return

    def disable(self):
        '''
        Disable trapping the break.  You can check whether a break
        was trapped using the count and trapped properties.
        '''
        if self._enabled:
            self._enabled = False
            signal.signal(signal.SIGINT, self._oldhandler)
            self._oldhandler = None
        return

    def __call__(self, signame, sf):
        '''
        An break just occurred.  Save information about it and keep
        going.
        '''
        self._count += 1
        # If we've exceeded the "emphatic" count disable this handler.
        if self._count >= self._emphatic:
            self.disable()
        return

    def __del__(self):
        '''
        Python is reclaiming this object, so make sure we are disabled.
        '''
        try:
            # On Linux, this might throw the following:
            # Exception SystemError: 'error return without exception set' in
            # <bound method BreakHandler.__del__ of
            # <__main__.BreakHandler instance at 0x7f16b6e30440>> ignored
            self.disable()
        except SystemError:
            pass
        return

    @property
    def count(self):
        '''
        The number of breaks trapped.
        '''
        return self._count

    @property
    def trapped(self):
        '''
        Whether a break was trapped.
        '''
        return self._count > 0


# Copied from http://stackoverflow.com/questions/11150239/python-natural-sorting
def natural_sort(l):
    convert = lambda text: int(text) if text.isdigit() else text.lower()
    alphanum_key = lambda key: [convert(c) for c in re.split('([0-9]+)', key['name'])]
    return sorted(l, key=alphanum_key)


def hostname():
    """Return the short hostname of the machine."""
    return (subprocess.Popen(["hostname", "-s"], stdout=subprocess.PIPE).
            communicate()[0].strip())


def is_integer(s):
    """Returns whether a string contains an integer."""
    return re.match(r'^\d+$', s) is not None


def get_start_command(node, shell):

    if shell == 'erl':
        erl_program = re.sub(r'(beam|beam.smp)$', 'erl', node['program'])
    elif shell == 'iex':
        erl_program = 'iex'

    client_name = ('connect2beam_%s_%s' %
                   (random.randint(0,10000), node['name']))

    # The node can be started in 3 ways:
    # * -name [email protected]   # long name
    # * -sname namebase@myhost     # short name 1
    # * -sname namebase            # short name 2
    # In the last case, we need to append the host name so that we
    # can connecto the machine.
    node_name = node['name']
    name_type = node['name_type']
    if name_type == '-sname' and '@' not in node_name:
        node_name += '@' + hostname()

    # Elixir has '--name' and '--sname' instead of '-name' and '-sname'.
    if shell == 'iex':
        name_type = '-' + name_type

    cookie = node['cookie']
    if cookie is None:
        cookie_opts = []
    else:
        if shell == 'erl':
            cookie_opts = ['-setcookie', cookie]
        elif shell == 'iex':
            cookie_opts = ['--cookie', cookie]

    start_cmd = [
        erl_program,
        name_type,
        client_name,
        '-remsh' if shell == 'erl' else '--remsh',
        node_name,
        '-hidden'] + cookie_opts

    return start_cmd


def connect_to_node(node, shell):
    # shell = 'erl' | 'iex'
    bh = BreakHandler()
    bh.enable()
    start_cmd = get_start_command(node, shell)
    print start_cmd
    p = subprocess.Popen(start_cmd).wait()


def main():

    out = subprocess.Popen(['ps', 'ax', '-o', 'pid,command'],
                           stdout=subprocess.PIPE).communicate()[0]
    processes_raw = out.split('\n')[1:-1]
    node_index = 1
    node_noconn = []
    node_conn = []

    for process_raw in processes_raw:
        process = process_raw.strip().split()
        pid = process[0]
        program = process[1]
        params = process[2:]
        if (program.endswith('beam') or program.endswith('beam.smp')):
            name = None
            name_type = None # '-name' or '-sname'
            cookie = None
            i = 0
            while i < len(params):
                param = params[i]
                if param in ('-name', '-sname'):
                    name = params[i + 1]
                    name_type = param
                    i += 1
                elif param == '-setcookie':
                    cookie = params[i + 1]
                    i += 1
                i += 1
            if None in (name, name_type):
                node_noconn.append(process_raw)
            else:
                node_conn.append({
                                  'name': name,
                                  'cookie': cookie,
                                  'name_type': name_type, # -name or -sname

                                  # The program that started this node
                                  'program': program,

                                  'pid': pid})

    node_conn = natural_sort(node_conn)
    if len(node_conn) > 0:
        print 'Available nodes: index name (cookie, pid)'
        print
        for i, node in enumerate(node_conn):
             print '    %s. %s (%s %s)' % (i + 1, node['name'], node['cookie'],
                                            node['pid'])

    if len(node_noconn) > 0:
        print '''
Cannot connect to the following nodes because node name is missing:'''
    for process in node_noconn:
        print ' ', process

    if len(node_conn) > 0:
        sys.stdout.write('''
Type either of the followings:
- The index of the node name to which you wish to connect.
- The letter "x" and the index of the node name to which you wish to connect
  using the iex shell.
- The letter "k" and the index of nodes you wish to kill (separated with
  spaces or with an 1..5 syntax).
- Hit Enter to cancel.

> ''')
        sys.stdout.flush()
    else:
        print 'No Erlang nodes found.'

    if len(node_conn) > 0:

        line = sys.stdin.readline().strip()
        if line == '':
            return
        words = line.split()

        if is_integer(words[0]):
            node_index = int(words[0])
            node = node_conn[node_index - 1]
            connect_to_node(node, 'erl')

        elif words[0] == 'x':
            node_index = int(words[1])
            node = node_conn[node_index - 1]
            connect_to_node(node, 'iex')

        elif words[0] == 'k':
            nodes_to_kill = []
            for word in words[1:]:
                if is_integer(word):
                    nodes_to_kill.append(int(word))
                else:
                    r = re.match(r'(\d+)\.\.(\d+)', word)
                    if r:
                        first_index = int(r.group(1))
                        last_index = int(r.group(2))
                        for node_index in range(first_index, last_index + 1):
                            nodes_to_kill.append(node_index)
                    else:
                        print 'Unknown format:', word
                        return
            for node_index in nodes_to_kill:
                node = node_conn[node_index - 1]
                cmd = ['kill', str(node['pid'])]
                print cmd
                p = subprocess.Popen(cmd).wait()

        else:
            print 'Unknown command:', words[0]


if __name__ == '__main__':
    main()