manage batch of server

q-goblin.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import multiprocessing
import paramiko
import getpass
import logging
import requests
import argparse
import socket
import time


class KeyboardInterruptError(Exception): pass


def ssh_connection(user, host, port=22, password=None, key_filename=None):
    """
    with password='' you will be prompted for a password when the script runs
    """
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    if password == '':
        # ask for one on stdin
        password = getpass.getpass('Password for %s@%s: ' % (user, host))
    try:
        ssh.connect(host, port=port, username=user, password=password, key_filename=key_filename, timeout=2.0)
    except Exception as e:
        time.sleep(1)
        ssh.connect(host, port=port, username=user, password=password, key_filename=key_filename, timeout=2.0)
        print e
        pass
    # custom attributes

    ssh.user = user
    if user == 'root':
        ssh.homedir = '/root'
    else:
        ssh.homedir = '/home/%s' % user
    ssh.host = host
    ssh.password = password
    ssh.use_sudo = False
    return ssh


def run_remote(ssh, cmd, check_exit_status=True, verbose=True):
    chan = ssh.get_transport().open_session()
    stdin = chan.makefile('wb')
    stdout = chan.makefile('rb')
    stderr = chan.makefile_stderr('rb')
    processed_cmd = cmd
    if ssh.use_sudo:
        processed_cmd = 'sudo -S bash -c "%s"' % cmd.replace('"', '\\"')
    chan.exec_command(processed_cmd)
    # If stdout is still open then sudo is asking us for a password
    if stdout.channel.closed is False:
        stdin.write('%s\n' % ssh.password)
        stdin.flush()
    result = {
        'stdout': [],
        'stderr': [],
    }
    exit_status = chan.recv_exit_status()
    result['exit_status'] = exit_status

    def get_output():
        for line in stdout:
            result['stdout'].append(line)
        for line in stderr:
            result['stderr'].append(line)
        return result
    outputs = []
    if check_exit_status and exit_status != 0:
        outputs = get_output()
    if verbose:
        outputs = get_output()
    return outputs


def start_process():
    pass
    #print 'Starting', multiprocessing.current_process().name


def do_process(operate, hostname, username, password):
    print "processing on %s ..." % print_blue(hostname)
    try:
        sudoconn = ssh_connection(username, hostname, password=password)
        sudoconn.use_sudo = True
        sudoconn.su_user = 'root'
        return globals()[operate](hostname, sudoconn)
    except KeyboardInterrupt:
        raise KeyboardInterruptError()


def do_restart_nginx(hostname, sudoconn):
    status = run_remote(sudoconn, "/etc/init.d/nginx restart && pidof nginx")
    if status['exit_status'] != 0:
        print("restart nginx [ %s ] %s" % (print_blue(hostname), print_red("failed")))
        return False
    else:
        print("restart nginx [ %s ] %s" % (print_blue(hostname), print_green("sucess")))


def do_reload_nginx(hostname, sudoconn):
    status = run_remote(sudoconn, "/etc/init.d/nginx reload")
    if status['exit_status'] != 0:
        print("reload nginx [ %s ] %s" % (print_blue(hostname), print_red("failed")))
        return False
    else:
        print("reload nginx [ %s ] %s" % (print_blue(hostname), print_green("sucess")))


def do_reload_php_55(hostname, sudoconn):
    status = run_remote(sudoconn, "/etc/init.d/php-fpm-5.5.0 reload")
    if status['exit_status'] != 0:
        print("reload php-5.5.0 [ %s ] %s" % (print_blue(hostname), print_red("failed")))
        return False
    else:
        print("reload php-5.5.0 [ %s ] %s" % (print_blue(hostname), print_green("sucess")))


def do_restart_php_55(hostname, sudoconn):
    status = run_remote(sudoconn, "/etc/init.d/php-fpm-5.5.0 restart")
    if status['exit_status'] != 0:
        print("reload php-5.5.0 [ %s ] %s" % (print_blue(hostname), print_red("failed")))
        return False
    else:
        print("reload php-5.5.0 [ %s ] %s" % (print_blue(hostname), print_green("sucess")))


def do_get_uptime(hostname, sudoconn):
    infos = run_remote(sudoconn, "uptime")
    for info in infos['stdout']:
        uptime = info.strip()
    print("[ %s ] %s" % (print_blue(hostname), print_green(uptime)))


def print_green(text):
    color_text = "\033[1;36m%s\033[0m" % text
    return color_text


def print_red(text):
    color_text = "\033[1;31m%s\033[0m" % text
    return color_text


def print_blue(text):
    color_text = "\033[1;34m%s\033[0m" % text
    return color_text


def get_hostnames(domain, host, idc):
    if host:
        return host
    return get_hosts_by_domain(domain, idc)


def get_hosts_by_domain(domain, idc=""):
    mapping = domain_mapping()
    if domain not in mapping:
        return False
    res = requests.get('http://api.xxx.com?domain=%s&port=80' % mapping[domain])
    hosts = []
    for host in res.json():
        if str(host).find(idc) >= 0:
            hosts.append(host)
    return hosts


def domain_mapping():
    mapping = {}
    mapping['aaa'] = 'aaa.com'
    mapping['bbb']   = 'bbb.com'
    mapping['ccc']    = 'ccc.com'
    return mapping


def get_domains(domains):
    print "available domains:"
    for k, v in domains.items():
        print("%s -- %s" % (print_green(k), v))


def get_hosts(domain):
    mapping = domain_mapping()
    if domain not in mapping:
        print "invalidate argument of domain, available domains:[ %s ]" % ", ".join(mapping.keys())
        return False
    for host in get_hosts_by_domain(domain):
        print(print_blue(host))


def operation_mapping():
    mapping = {}
    mapping['restart-nginx'] = "do_restart_nginx"
    mapping['reload-nginx']  = "do_reload_nginx"
    mapping['restart-php5.5'] = "do_restart_php_55"
    mapping['reload-php5.5']  = "do_reload_php_55"
    mapping['get-status']    = "do_get_uptime"
    return mapping


def get_operation(operation):
    if operation == "":
        return False
    mapping = operation_mapping()
    if operation not in mapping:
        return False
    return mapping[operation]


if __name__ == '__main__':
    parser = argparse.ArgumentParser("^_^")
    #parser = parser.add_mutually_exclusive_parser()
    parser.add_argument("--idc", action='store', dest='_idc', default="", choices=['zwt', 'bjsc', 'shm'], help="set idc[zwt,bjsc,shm] of servers")
    parser.add_argument("--host", action='append', dest='_hosts', default=[], help="set host or ip of servers")
    parser.add_argument("--domain", action='store', dest='_domain', default="", help="set domain")
    parser.add_argument("--operation", action='store', dest='_operation', default="", help="set the operation todo with servers")
    parser.add_argument("--hosts", action='store_true', dest='_list_hosts', default=False, help="list hosts by domain")
    parser.add_argument("--domains", action='store_true', default=False, dest='_list_domains', help="list available domains")
    args = parser.parse_args()
    if args._list_domains is True:
        get_domains(domain_mapping())
        exit()
    if args._list_hosts is True:
        get_hosts(args._domain)
        exit()
    operation = get_operation(args._operation)
    hostnames = get_hostnames(args._domain, args._hosts, args._idc)
    if operation is False or hostnames is False:
        print("at least one argument is required, Use [%s -h] for help" % __file__)
        exit()

    login_user = raw_input('Username: ')
    if login_user == '':
        login_user = 'xxxx'
    password = getpass.getpass('Password :')
    multiprocessing.log_to_stderr(logging.WARNING)
    pool_size = multiprocessing.cpu_count() * 2
    pool = multiprocessing.Pool(processes=pool_size,
                                initializer=start_process,
                                #maxtasksperchild=2,
                                )
    try:
        for host in hostnames:
            outputs = pool.apply(do_process, (operation, host, login_user, password,))
            to_continue = 'y'
            if outputs is False:
                to_continue = raw_input('Process failed [ %s ]: [y/n]' % host)
                if to_continue.lower() == 'y':
                    continue
                else:
                    break
        pool.close()
    except KeyboardInterrupt:
        print 'got ^C while pool mapping, terminating the pool'
        pool.terminate()
    except Exception, e:
        print 'got exception: %r, terminating the pool' % (e,)
        pool.terminat()
    finally:
        pool.join()
        print 'process complete'