ceiling.py from EmStar at Krugle

Show ceiling.py syntax highlighted

#!/usr/bin/python

import os
import sys
import time
import getopt
import shutil
import signal
import getpass
import commands

default_cfg_file = "stargate-mappings.cfg"
temp_image_dir = "temp_mote_images_" + getpass.getuser()
ceiling_keys_location = "http://lecs.cs.ucla.edu/~mlukac/ceiling-keys.tar.gz"
binary_deps = ["wget", "avr-objdump", "set-mote-id", "scp"]


def usage():
    print "./ceiling.py -i|--image moteimage -e|--emstar EmStarRoot [-c|--config configfile]"
    print "-i|--image moteimage (required)"
    print "-e|--emstar EmStarRoot (required)"
    print "-c|--config configfile (optional... looks for stargate-mappings.cfg)"
    print "Note: I will remove than create a directory for the temporary images called", temp_image_dir
    sys.exit(1)


def check_path():
    global binary_deps
    missing = 0
    allpaths = os.path.expandvars("$PATH").split(":")
    for bin in binary_deps:
        found = 0
        for path in allpaths:
            if os.path.exists(os.path.join(path, bin)):
                found = 1
                break
        if found == 0:
            print "Can not find " + bin + " in path"
            missing = missing + 1
    if missing > 0:
        print "Couldn't find " + str(missing) + " of the " + str(len(binary_deps)) + " required programs"
        sys.exit(1)


def read_in_cfg(cfg_file):
    output_map = {}
    f = open(cfg_file)
    lines = f.readlines()
    for line in lines:
        if line.startswith("#") or line.startswith("\n"):
            continue
        map = line.split(' ')
        output_map[map[0]] = map[1].strip()
    return output_map


def create_mote_images(node_mappings, mote_image):
    # setup an output dir
    if os.path.exists(temp_image_dir):
        shutil.rmtree(temp_image_dir)
    os.mkdir(temp_image_dir)
    shutil.copy(mote_image, mote_image + ".srec")
    for node in node_mappings.keys():
        exec_string = "%s --srec %s.srec %s/%s.%s %s" % ("set-mote-id", mote_image,
                                                        temp_image_dir, mote_image,
                                                        node, node)
        #print "Setting mote id with command"
        #print exec_string
        if os.system(exec_string) != 0:
            print "Error creating mote image", node


def setup_keys():
    print "I'm temporarily moving your ssh keys to ~/.ssh-orig"
    if os.path.isfile(os.path.expanduser("~/.burningmotes")):
        return
    if os.path.isdir(os.path.expanduser("~/.ssh-backup-moteburn")) is False:
        os.system("cp -r ~/.ssh ~/.ssh-backup-moteburn")
    if os.path.isdir(os.path.expanduser("~/.ssh-ceiling")):
        print "Keys exists... not downloading"
    else:
        os.system("wget " + ceiling_keys_location)
        os.system("tar -C ~/ -zxf ceiling-keys.tar.gz")

    if os.path.isdir(os.path.expanduser("~/.ssh")):
        print "backing up your ssh keys"
        os.system("mv ~/.ssh ~/.ssh-orig")
    os.system("rm -rf ~/.ssh")
    os.system("mv ~/.ssh-ceiling ~/.ssh")
    os.system("rm -rf ceiling-keys.tar.gz")
    os.system("touch ~/.burningmotes")


def unsetup_keys():
    print "Returning your ssh keys"
    os.system("mv ~/.ssh ~/.ssh-ceiling")
    if os.path.isdir(os.path.expanduser("~/.ssh-orig")):
        os.system("mv ~/.ssh-orig ~/.ssh")
    os.system("rm -rf ~/.burningmotes")


def copy_files(image_name, id_to_name):
    for id in id_to_name.keys():
        # do some thread magic in future versions
        print "Copying to " + id_to_name[id]
        ssh_cmd = "scp -B %s/%s.%s root@%s:/tmp/" % (temp_image_dir, image_name, id,
                                                     id_to_name[id])
        status, output = commands.getstatusoutput(ssh_cmd)
        print "Exit status: " + str(status)
        print output


def start_fusdnet(id_to_name, pids, fusdnet_client):
    for id in id_to_name.keys():
        try:
            pid = os.fork()
            if pid > 0:
                print "Starting fusdnet for " + str(id) + " as pid " + str(pid)
                pids.append((pid, id))
            else: # run fusddnet
                os.execlp(fusdnet_client,
                       fusdnet_client,
                       "fusd://" + id_to_name[id] +
                          "/mote/0/program=/dev/mote/" + id_to_name[id] + "/program")
                sys.exit(1)
        except OSError, e:
            print >>sys.stderr, "fork failed: %d (%s)" % (e.errno, e.strerror)
            sys.exit(1)


def kill_fusdnet(pids):
    for pid, id in pids:
        print "killing " + str(id) + "s fusdnet pid " + str(pid)
        os.kill(pid, signal.SIGHUP)
    # maybe wait? or wait when rebooting stargates


def run_tosinstall(id_to_name, image_name, tosinstall_bin):
    global emstar_root
    pids = {}
    count = 0
    for id in id_to_name.keys():
        try:
            pid = os.fork()
            if pid > 0:
                print "Starting tos install for " + str(id) + " as pid " + str(pid)
                pids[pid] = id;
                count = count + 1
            else:
                os.execlp(tosinstall_bin,
                          tosinstall_bin,
                          "--prog", "stargate-remote",
                          "--root", emstar_root,
                          "--rroot", "/mnt/nfs/girod/emstar",
                          "--index", id_to_name[id],
                          "/tmp/" + image_name + "." + id)
                sys.exit(1)
        except OSError, e:
            print >>sys.stderr, "fork failed: %d (%s)" % (e.errno, e.strerror)
            sys.exit(1)
    # wait for them to finish
    # !!! THIS probably won't work... especially if one of them exists before I spawn them all !!!
    print "!!! Waiting for burns to complete !!!"
    print "!!! This will take a while !!!"
    while count > 0:
        pid, status = os.waitpid(-1,0)
        print "pid " + str(pid) + " node " + pids[pid] + " exited with status " + str(status)
        count = count - 1


def check_tosinstall(id_to_name):
    print "Printing status from install in 2 seconds"
    print "Did any of these fail? ... if So, you probbaly want check it out by hand"
    time.sleep(2)
    for name in id_to_name.values():
        status, output = commands.getstatusoutput("cat /dev/mote/" + name + "/program")
        print output
        

def reboot_stargates(id_to_name):
    for id in id_to_name.keys():
        print "Rebooting " + id_to_name[id]
        ssh_cmd = "ssh root@%s \"shutdown -r now\"" % (id_to_name[id])
        status, output = commands.getstatusoutput(ssh_cmd)
        print "Exit status: " + str(status)
        print output


#######################
# Do the option stuff #
#######################

# options = ( opts, remaining_args )
# opts = list of ( opt, param )
options = getopt.getopt(sys.argv[1:], "i:c:e:", ["image=", "config=", "emstar="])

if len(options[1]) > 0:
    usage()

cfg_file = default_cfg_file
image_file = ""
emstar_root = ""

for opt in options[0]:
    if opt[0] == "-i" or opt[0] == "--image":
        image_file = opt[1]
    elif opt[0] == "-c" or opt[0] == "--config":
        cfg_file = opt[1]
    elif opt[0] == "-e" or opt[0] == "--emstar":
        emstar_root = opt[1]
    else:
        usage()

if image_file == "" or os.path.isfile(image_file) is False:
    print "Please provide valide mote image"
    usage()

if emstar_root == "" or os.path.isdir(emstar_root) is False:
    print "Please provide a valid emstar root"

if os.path.isfile(cfg_file) is False:
    print "Please provide valid cfg file, or create default stargate-mappings.cfg"
    usage()

fusdnet_client_bin = os.path.join(emstar_root, "obj.i686-linux/fusd/fusdnet_client")
tosinstall_bin = os.path.join(emstar_root, "bin/tosinstall")
if os.path.isfile(fusdnet_client_bin) is False:
    print "Can't find fusdnet client at : " + fusdnet_client_bin
    usage()

if os.path.isfile(tosinstall_bin) is False:
    print "Can't find tosinstall at: " + tosinstall_bin
    usage()
        

###########
# main(): #
###########

# check to see if the requried programs are in the path
check_path()

# read in the cfg mapping file
node_stargate_mapping = read_in_cfg(cfg_file)

# create the mote images with the right ID's
create_mote_images(node_stargate_mapping, image_file)

# setup the users ssh keys so they can ssh to the stargates
setup_keys()

# scp the files to the stargates
copy_files(image_file, node_stargate_mapping)

# startup the fusdnet_client's 
fusdnet_pids = []
start_fusdnet(node_stargate_mapping, fusdnet_pids, fusdnet_client_bin)
print "sleeping for 2 seconds ... "
time.sleep(2)
#time.sleep(10)

# tell the stargates to burn. This will block
run_tosinstall(node_stargate_mapping, image_file, tosinstall_bin)

# double check
check_tosinstall(node_stargate_mapping)

# reboot the stargates
reboot_stargates(node_stargate_mapping)

# kill all the fusdnet clients
kill_fusdnet(fusdnet_pids)

# put the users keys back, so they won't notice a thing
unsetup_keys()
See more files for this project here