digitalocean_config.py - setupup digitalocean droplets to use python plumbum

digitalocean_config.py

from plumbum import SshMachine
from plumbum import local
from plumbum.cmd import sudo
from plumbum.cmd import grep, wc, cat, head
import os,subprocess,digitalocean,shutil,glob,os.path,plumbum
import servers.secret as secret
from servers.common import sshDir,create_ssh_dir
#http://www.robotgoblin.co.uk/blog/2012/07/24/managing-multiple-ssh-keys/

def read_hosts():
    hostPath = os.path.join(os.environ['HOME'],"hosts.orig")
    if not os.path.exists(hostPath):
        shutil.copy("/etc/hosts",hostPath)
    
    with open(hostPath,"rt") as f:
        return map(lambda l : l.strip(),f.readlines())
        
#rem = SshMachine("128.199.177.1", user = "root", keyfile = "/home/john/.ssh/id_rsa")
def read_idRsaPub():
    with open(os.path.join(sshDir,"docker_id_rsa.pub"),"rb") as f:
        return f.read()
    
def read_idRsa():
    with open(os.path.join(sshDir,"docker_id_rsa"),"rb") as f:
        return f.read()
    
def create_hosts_file(dockerContainer = None):
    create_ssh_dir()
    hostFile = read_hosts()
    knownHostsPath = os.path.join(sshDir,"known_hosts")
    tmpHostPath = "/tmp/hosts"
    try:
        os.remove(knownHostsPath)
    except OSError:
        pass
    for droplet in droplets:
        print(droplet.ip_address,droplet.name)
        hostFile.append("%s %s" % (droplet.ip_address,droplet.name))
        
    if dockerContainer :
        print("docker container host file")
        info = docker.inspect_container(dockerContainer)
        hostsPath = info["HostsPath"]
        hostBash = []
        hostBash.append("#!/bin/bash")
        hostBash.append("sudo mv %s %s" % ("/root/hosts",hostsPath))
        with open("/Development/copy_hosts.sh","wt") as f:
            f.write("\n".join(hostBash))
        with open("/Development/hosts","wt") as f:
            f.write("\n".join(hostFile))
    else:
        with open(tmpHostPath,"wt") as f:
            f.write("\n".join(hostFile))
        os.system('sudo mv /tmp/hosts /etc/hosts')
    
def add_to_known_hosts(droplets):
    for droplet in droplets:
        ipScanStr = "ssh-keyscan -H {hostName} >> {sshDir}/known_hosts"
        nameScanStr = "ssh-keyscan {hostName} >> {sshDir}/known_hosts"
        subprocess.call(ipScanStr.format(hostName = droplet.ip_address,sshDir = sshDir),shell=True)
        subprocess.call(nameScanStr.format(hostName = droplet.name,sshDir = sshDir),shell=True)

configTemplate = """
Host {ipAddress}
    User root
    Hostname {ipAddress}
    PreferredAuthentications publickey
    IdentityFile ~/.ssh/d_{hostName}/id_rsa
"""
    
def delete_docker_dirs():
    create_ssh_dir()
    fileNames = os.listdir(sshDir)
    for fileName in fileNames:
        if fileName.startswith("d_"):
            fileName = os.path.join(sshDir,fileName)
            if os.path.isdir(fileName):
                shutil.rmtree(fileName)

def create_ssh_config(droplets):
    create_ssh_dir()
    delete_docker_dirs()
    configFile = []
    for droplet in droplets:
        kDir = os.path.join(sshDir,"d_{hostName}".format(hostName = droplet.name))
        os.mkdir(kDir)
        print(droplet.ip_address,droplet.name)
        configFile.append(configTemplate.format(hostName = droplet.name,ipAddress = droplet.ip_address))
        with open(os.path.join(kDir,"id_rsa"),"wt") as f:
            f.write(secret.idRsa)
        subprocess.call("sudo chmod 600 %s" % os.path.join(kDir,"id_rsa"),shell=True)
        
    with open(os.path.join(sshDir,"config"),"wt") as f:
        f.write("\n".join(configFile))

def display_droplets(droplets):
    import ipython_helpers
    keys = ['index',"name","ip_address","created_at","image","size"]
    def merge(i,droplet):
        result = {"index" : i}
        result.update(droplet.__dict__)
        return result
    return ipython_helpers.dicts_to_html([merge(i,droplet) for i,droplet in zip(range(100),droplets)],keys = keys)

def display_images(manager):
    import ipython_helpers
    images = manager.get_my_images()
    keys = ['index','name','type','created_at','min_size','min_disk_size']
    def merge(i,image):
        result = {"index" : i}
        result.update(image.__dict__)
        return result
    return ipython_helpers.dicts_to_html([merge(i,image) for i,image in zip(range(100),images)],keys = keys)

def display_sizes(manager):
    """
    dict_keys(['token', 'available', 'vcpus', 'transfer', 'price_hourly', 'slug', 'regions', '_log', 'price_monthly', 'memory', 'disk', 'end_point'])
    """
    import ipython_helpers
    def merge(i,slug):
        result = {"index" : i}
        result.update(slug.__dict__)
        return result
    
    keys = ["index",'price_monthly','price_hourly','vcpus','memory', 'disk']
    slugSizes = manager.get_all_sizes()
    return ipython_helpers.dicts_to_html([merge(i,slug) for i,slug in zip(range(100),slugSizes)],keys = keys)


def create_ssh(droplets,user = "root",keyfile = "/home/john/docker_id_rsa"):
    sshMachines = []
    for droplet in droplets:
        sshMachine = SshMachine(droplet.ip_address,user = "root", keyfile = keyfile)
        sshMachines.append(sshMachine)
    return sshMachines

def create_ssh(droplets,user = "root",keyfile = "/home/john/docker_id_rsa"):
    sshMachines = []
    for droplet in droplets:
        sshMachine = SshMachine(droplet.ip_address,user = "root", keyfile = keyfile)
        sshMachines.append(sshMachine)
    return sshMachines