Add script to test if a testnet is ready

#!/usr/bin/env python3

from argparse import RawTextHelpFormatter, ArgumentParser

from stem.control import Controller

import stem

import time

import os

import logging

logger = logging.getLogger(__name__)

def get_controller(sock_fname):

    cont = Controller.from_socket_file(path=sock_fname)

    cont.authenticate()

    return cont

def get_is_bootstrapped(cont, timeout=60):

    start_time = time.time()

    while start_time + timeout > time.time():

        line = cont.get_info('status/bootstrap-phase')

        state, _, progress, *_ = line.split()

        progress = int(progress.split('=')[1])

        if state == 'NOTICE' and progress == 100:

            logger.debug('Tor is bootstrapped')

            return True

        time.sleep(1)

    logger.debug('Tor didn\'t bootstrap before timeout')

    return False

def get_has_full_consensus(cont, network_size, timeout=60):

    start_time = time.time()

    while start_time + timeout > time.time():

        relays = [r for r in cont.get_network_statuses()]

        if len(relays) == network_size:

            logger.debug('Tor has correct network size %d',

                         network_size)

            return True

        elif len(relays) > network_size:

            logger.warning('Tor has more relays than expected. %d vs %d',

                           len(relays), network_size)

            return True

        time.sleep(1)

    logger.debug('Tor didn\'t reach expected network size %d before '

                 'timeout', network_size)

    return False

def is_relay_ready(sock_fname, network_size):

    with get_controller(sock_fname) as cont:

        if not get_is_bootstrapped(cont):

            logger.warning('%s not bootstrapped, Tor not ready', sock_fname)

            return False

        if not get_has_full_consensus(cont, network_size):

            logger.warning('%s doesn\'t have full consensus, Tor not ready',

                           sock_fname)

            return False

    logger.info('%s is ready', sock_fname)

    return True

def main(args):

    for datadir in args.datadir:

        logger.info('Checking %s', datadir)

        sock_fname = os.path.join(datadir, 'control_socket')

        assert os.path.exists(sock_fname)

        if not is_relay_ready(sock_fname, network_size=args.size):

            return 1

    # If we got to this point, it seems like every relay is completely ready.

    # Do one more check to make sure that's still the case.

    for datadir in args.datadir:

        logger.info('Verifying %s', datadir)

        sock_fname = os.path.join(datadir, 'control_socket')

        assert os.path.exists(sock_fname)

        if not is_relay_ready(sock_fname, network_size=args.size):

            return 1

    return 0

if __name__=='__main__':

    desc = '''

Given the data directories for a local tor network, connect to the control

socket in each directory and verify that the tor on the other end of the socket

is fully bootstrapped and has the right size of consensus.

The "right size of consensus" is determined based on the number of data

directories given to check. If that is not okay to assume (for example, there

are some client data directories given to check), then specify the size

manually with --size.

Waits up to 60 seconds for each check for each tor.

- In the worst case, this script will take a long time to run (if every tor

  suddenly passes each check after 59 seconds).

- In the normal failure case, this script will take about 60 seconds to run

  (the first tor is not ready and fails its checks).

- In the normal case, it will run very quickly (every tor is bootstrapped and

  ready).

Exits with 0 if everything is good. Otherwise exits with a postive integer.

'''

    parser = ArgumentParser(

            formatter_class=RawTextHelpFormatter, description=desc)

    parser.add_argument('-s', '--size', type=int, help='If given, don\'t '

                        'assume the network size based on the number of '

                        'datadirs, but use this size instead.')

    parser.add_argument('-d', '--debug', action='store_true')

    parser.add_argument('datadir', nargs='+', type=str)

    args = parser.parse_args()

    if args.debug:

        logger.setLevel(logging.DEBUG)

    else:

        logger.setLevel(logging.WARNING)

    ch = logging.StreamHandler()

    formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')

    ch.setFormatter(formatter)

    logger.addHandler(ch)

    if not args.size:

        args.size = len(args.datadir)

    try:

        exit(main(args))

    except KeyboardInterrupt:

        pass