uow

#!/usr/bin/env python
import string
import optparse
import os.path as osp
import os
import sys
import shutil
import subprocess
import datetime
import ConfigParser as configparser

NGINX_SERVER_TMPL = '''\
# Generated nginx virtual host configuration. 
# DO NOT EDIT MANUALLY. Use the uow script to update this file.
${upstreams}

${http_insert}

server {
    listen ${server_port};
    server_name ${server_name};
    server_name_in_redirect off;
    location /libs {
        alias ../servers/${config_name}/www/libs;
    }
${aliases}
${upload}
${proxies}
${server_insert}
}'''
NGINX_UPSTREAM_TMPL = '''\
upstream ${name} {
    ${servers}
}'''
NGINX_ALIAS_TMPL = '''\
    location ${url} {
        alias ${path};
    }
'''
NGINX_PROXY_TMPL = '''\
    location ${url} {
        proxy_pass_header Server;
        proxy_set_header Host $host;
        proxy_redirect off;
        proxy_set_header X-Real-IP $remote_addr;
        proxy_set_header X-Scheme $scheme;
        proxy_pass ${upstream};
    }'''
    
NGINX_UPLOAD_TMPL = '''\
    location ${url} {
        # Pass altered request body to this location
        upload_pass ${upstream};

        # Store files to this directory
        # The directory is hashed, subdirectories 0 1 2 3 4 5 6 7 8 9 should exist
        upload_store /var/tmp/uow;
        
        # Set specified fields in request body
        upload_set_form_field $upload_field_name.name "$upload_file_name";
        upload_set_form_field $upload_field_name.content_type "$upload_content_type";
        upload_set_form_field $upload_field_name.path "$upload_tmp_path";

        # Inform backend about hash and size of a file
        #upload_aggregate_form_field "$upload_field_name.md5" "$upload_file_md5";
        #upload_aggregate_form_field "$upload_field_name.size" "$upload_file_size";

        upload_pass_form_field "^tags$|^Authorization$|^title$|^description$|^creditURL$";

        upload_cleanup 400 404 499 500-505;
    }'''


class NginxTemplate(object):
    def __init__(self, config_name):
        self.template = string.Template(NGINX_SERVER_TMPL)
        self.config_name = config_name
        self.upstreams = []
        self.proxies = []
        self.upload = []
        self.aliases = []
        self.server_name = None
        self.server_port = None
        self.server_insert = ''
        self.http_insert = ''
    
    def render(self):
        args = dict(self.__dict__)
        args['upstreams'] = '\n'.join(self.upstreams)
        args['aliases'] = '\n'.join(self.aliases)
        args['proxies'] = '\n'.join(self.proxies)
        args['upload'] = '\n'.join(self.upload)
        return self.template.substitute(args)

    def assert_complete(self):
        assert(self.server_port is not None and self.server_name is not None)
    
    def add_alias(self, url, path):
        config_name = self.config_name
        if not path.startswith('/'):
            path = os.path.join('../servers', config_name, path)
        tmpl = string.Template(NGINX_ALIAS_TMPL)
        text = tmpl.safe_substitute(locals())
        self.aliases.append(text)
    
    def add_proxy(self, url, upstream):
        tmpl = string.Template(NGINX_PROXY_TMPL)
        text = tmpl.safe_substitute(locals())
        self.proxies.append(text)
        
    def add_upload(self, url, upstream):
        tmpl = string.Template(NGINX_UPLOAD_TMPL)
        text = tmpl.safe_substitute(locals())
        self.upload.append(text)
        
    def add_upstream(self, name, servers):
        servers = '\n'.join(['server %s;' % server for server in servers])
        tmpl = string.Template(NGINX_UPSTREAM_TMPL)
        text = tmpl.substitute(locals())
        self.upstreams.append(text)
        
    def set_server_port(self, port):
        self.server_port = port
        
    def set_server_name(self, name):
        self.server_name = name
        
    def set_server_insert(self, value):
        self.server_insert = value
        
    def set_http_insert(self, value):
        self.http_insert = value

class NginxHandler(object):
    def __init__(self, config_name, nginx_tmpl):
        self.config_name = config_name
        self.nginx_tmpl = nginx_tmpl

    def dispatch(self, section, name, value, cp):
        getattr(self, name)(section, value, cp)
        
    def server_port(self, section, value, cp):
        self.nginx_tmpl.set_server_port(int(value))
    
    def server_name(self, section, value, cp):
        self.nginx_tmpl.set_server_name(value)
        
    def server_insert(self, section, value, cp):
        self.nginx_tmpl.set_server_insert(value)
        
    def http_insert(self, section, value, cp):
        self.nginx_tmpl.set_http_insert(value)

class ComponentHandler(object):
    def __init__(self, config_name, nginx_tmpl):
        self.config_name = config_name
        self.nginx_tmpl = nginx_tmpl
        
    def dispatch(self, section, name, value, cp):
        getattr(self, name)(section, value, cp)
    
    def update(self):
        pass

    def init(self, section, value, cp):
        os.system(value)

    def program(self, section, value, cp):
        # exist to avoid unsupported key warning
        pass
    
    def symlinks(self, section, value, cp):
        pairs = value.split(',')
        for source, target in zip(pairs[::2], pairs[1::2]):
            source, target = source.strip(), target.strip()
            dtarget = osp.dirname(target)
            # make source relative to target
            rel = osp.join(*(['..'] * len(dtarget.split('/'))))
            source = osp.join(rel, source)
            # build intervening paths
            try:
                os.makedirs(dtarget)
            except OSError:
                pass
            os.symlink(source, target)
            print 'Symlinked %s -> %s' % (target, source)
    
    def proxies(self, section, value, cp):
        program = cp.get(section, 'program')
        psection = 'program:%s' % program
        # create upstreams for proxies
        start = cp.getint(psection, 'first_port')
        count = cp.getint(psection, 'numprocs')
        servers = ['127.0.0.1:%d' % port for port in range(start, start+count)]
        self.nginx_tmpl.add_upstream(program, servers)
        # add proxies
        pairs = value.split(',')
        for url, path in zip(pairs[::2], pairs[1::2]):
            self.nginx_tmpl.add_proxy(url, path)
            print 'Proxied %s -> %s' % (path, url)
            
    def upload(self, section, value, cp):
        # add uploads
        pairs = value.split(',')
        for url, path in zip(pairs[::2], pairs[1::2]):
            self.nginx_tmpl.add_upload(url, path)
            print 'Upload %s -> %s' % (path, url)
            
    def aliases(self, section, value, cp):
        pairs = value.split(',')
        for url, path in zip(pairs[::2], pairs[1::2]):
            self.nginx_tmpl.add_alias(url, path)
            print 'Aliased %s -> %s' % (path, url)

class ProgramHandler(object):
    def __init__(self, config_name, nginx_tmpl):
        self.config_name = config_name
        self.nginx_tmpl = nginx_tmpl
        
    def dispatch(self, section, name, value, cp):
        getattr(self, name)(section, value, cp)
    
    def user(self, section, value, cp):
        os.seteuid(0)
        try:
            path = cp.get(section, 'directory')
            path = '/'.join(path.split('/')[2:])
            subprocess.check_call(['chown', '-h', value, path])
            subprocess.check_call(['chown', '-R', '-L', value, path])
        except subprocess.CalledProcessError, e:
            print 'ERROR: could not set ownership properly'
            raise e
        finally:
            os.seteuid(int(os.environ.get('SUDO_UID', 0)))
        print 'Set ownership %s -> %s' % (value, path)

def _prompt(msg, default):
    while 1:
        i = raw_input('%s [%s] ' % (msg, default)) or default
        i = i.lower()
        if i in ['y', 'yes']:
            return True
        elif i in ['n', 'no']:
            return False

def init(config_name, do_enable=True):
    # require root to set owner / group
    if os.getuid():
        raise ValueError('You must su or sudo to root to init.')
    # get pre-sudo user
    owner = int(os.environ.get('SUDO_UID', 0))
    # get path info
    cwd = os.getcwd()
    dname = osp.join('servers', config_name)
    cname = osp.join('servers', config_name+'.conf')
    # read config file
    cp = configparser.ConfigParser()
    cp.readfp(open(cname))
    
    # check if dir exists
    if osp.isdir(dname):
        # confirm dir reset
        reset = _prompt('Are your sure you want to delete %s?' % config_name, 'n')
        if reset:
            print 'Deleting server directories ...'
            shutil.rmtree(dname)
            print 'Deleting server nginx config ...'
            try:
                os.unlink(osp.join('nginx', 'servers', config_name+'.nginx.conf'))
            except OSError:
                pass
            print 'Deleting active symlink ...'
            try:
                os.unlink(osp.join('servers', 'active', config_name+'.conf'))
            except OSError:
                pass
            print 'Done deleting'
        else:
            print 'Aborted'
            return

    # switch to pre-sudo user
    os.seteuid(owner)
    # switch into dname to somewhat sandbox commands
    os.makedirs(dname)
    os.chdir(dname)
    
    # create standard folders
    print 'Creating server directories ...'
    # rwx, rx, rx
    os.makedirs(osp.join('www', 'libs'), 0765)
    os.makedirs(osp.join('services', 'logs'), 0765)
    os.makedirs(osp.join('components'), 0765)
    print 'Done creating server directories\n'
    # start an nginx template
    tmpl = NginxTemplate(config_name)

    # get nginx info
    nginx = NginxHandler(config_name, tmpl)
    print 'Handling "nginx" section ...'
    for name, value in cp.items('nginx'):
        try:
            # dispatch to handle keys
            nginx.dispatch('nginx', name, value, cp)
        except AttributeError:
            pass
            #print 'Unknown key "%s" in section "nginx"' % (name)
    print 'Done handling "nginx"\n'

    # get components
    components = ComponentHandler(config_name, tmpl)
    for section in filter(lambda x: x.startswith('component'), cp.sections()):
        print 'Handling "%s" section ...' % section
        if cp.has_option(section, 'init'):
            components.dispatch(section, 'init', cp.get(section, 'init'), cp)
        for name, value in cp.items(section):
            if name in ['init', 'backup']: continue
            try:
                # dispatch to handle keys
                components.dispatch(section, name, value, cp)
            except AttributeError:
                print 'Unknown key "%s" in section "%s"' % (name, section)
        print 'Done handling "%s"\n' % section

    # chown everything to the original user
    print 'Setting ownership to user %s' % owner
    os.chdir(cwd)
    os.seteuid(0)
    try:
        subprocess.check_call(['chown', '-R', '-L', str(owner), dname])
        subprocess.check_call(['chown', '-R', '-h', str(owner), dname])
    except subprocess.CalledProcessError, e:
        print 'ERROR: could not set ownership properly'
        raise e
    os.seteuid(owner)
    os.chdir(dname)
    print 'Done setting user\n'

    # get programs
    programs = ProgramHandler(config_name, tmpl)
    for section in filter(lambda x: x.startswith('program'), cp.sections()):
        print 'Handling "%s" section ...' % section
        if cp.has_option(section, 'user'):
            programs.dispatch(section, 'user', cp.get(section, 'user'), cp)
        print 'Done handling "%s"\n' % section

    # make sure we have all required nginx info
    tmpl.assert_complete()    
    # write nginx config
    conf = tmpl.render()
    os.chdir(cwd)
    fh = file(osp.join('nginx', 'servers', config_name+'.nginx.conf'), 'w')
    fh.write(conf)
    fh.close()
    # enable the config
    if do_enable:
        enable(config_name)
        
def _backup(config_name):
    # backup existing
    dt = datetime.datetime.now()
    bsrc = osp.join('servers', config_name)
    btarget = osp.join('backups', 
        config_name + '.' + dt.strftime('%Y-%m-%d.%H:%M:%S'))
    print 'Backing up %s to %s' % (bsrc, btarget)
    # use mv to retain owner/group; shutil.move loses it
    try:
        subprocess.check_call(['mv', bsrc, btarget])
    except subprocess.CalledProcessError, e:
        print 'ERROR: could not make backup copy'
        raise e
    return btarget

def update(config_name):
    # require root to set owner / group
    if os.getuid():
        raise ValueError('You must su or sudo to root to update.')

    # check if dir exists
    dname = osp.join('servers', config_name)
    if not osp.isdir(dname):
        print 'ERROR: no such server: %s' % config_name
        return

    # confirm dir reset
    reset = _prompt('Are your sure you want to update %s?' % config_name, 'n')
    if not reset:
        print 'Aborted'
        return

    # backup existing
    btarget = _backup(config_name)
    # disable existing
    try:
        disable(config_name, False)
    except IOError:
        pass
    # init existing
    try:
        init(config_name, False)
    finally:
        # init drops to pre-sudo user, set it back
        os.seteuid(0)
    # restore important sections
    cp = configparser.ConfigParser()
    cp.readfp(open(osp.join('servers', config_name + '.conf')))
    print 'Restoring backup ...'
    for section in filter(lambda x: x.startswith('component'), cp.sections()):
        try:
            rtarget = cp.get(section, 'backup')
        except configparser.NoOptionError:
            continue
        bsegs = btarget.split('/')
        rsegs = rtarget.split('/')
        rsrc = os.path.join(*(bsegs[:2] + rsegs[2:]))
        # remove dir created during init
        try:
            shutil.rmtree(rtarget)
        except OSError:
            pass
        # use system to preserve owner
        try:
            subprocess.check_call(['cp', '-a', rsrc, rtarget])
        except subprocess.CalledProcessError, e:
            print 'ERROR: could not restore from backup'
            raise e
        print 'Restored %s to %s' % (rsrc, rtarget)
    print 'Done restoring backup'
    # enable configuration again
    enable(config_name)

def remove(config_name):
    # require root to set owner / group
    if os.getuid():
        raise ValueError('You must su or sudo to root to remove.')
    # check if dir exists
    dname = osp.join('servers', config_name)
    if not osp.isdir(dname):
        print 'ERROR: no such server: %s' % config_name
        return

    # confirm dir reset
    reset = _prompt('Are your sure you want to delete %s?' % config_name, 'n')
    if not reset:
        print 'Aborted'
        return
    # create a backup
    _backup(config_name)
    try:
        os.unlink(osp.join('nginx', 'servers', config_name+'.nginx.conf'))
    except OSError:
        pass
    # disable the config
    try:
        disable(config_name, False)
    except IOError:
        # ignore if wasn't active
        pass

def enable(config_name):
    # check if dir exists
    dname = osp.join('servers', config_name)
    if not osp.isdir(dname):
        print 'ERROR: no such server: %s' % config_name
        return
    
    print 'Enabling "%s" config ...' % config_name
    target = osp.join('servers', 'active', config_name+'.conf')
    try:
        os.symlink(osp.join('..', config_name+'.conf'), target)
    except OSError:
        print 'WARNING: symlink already exists: %s' % config_name+'.conf'

    # read the config file
    cp = configparser.ConfigParser()
    cp.readfp(open(target))

    # read all the groups
    sections = (g for g in cp.sections() if g.startswith('group'))
    groups = (g.split(':')[1] for g in sections)
    
    # add and start the groups in supervisor
    try:
        subprocess.check_call(['supervisorctl', 'reread'])
    except subprocess.CalledProcessError, e:
        print 'ERROR: could not reread supervisord config'
        return
    for group in groups:
        print 'Adding and starting %s ...' % group
        try:
            subprocess.check_call(['supervisorctl', 'add', group])
            subprocess.check_call(['supervisorctl', 'start', '%s:*' % group])
        except subprocess.CalledProcessError, e:
            print 'ERROR: could not add / start supervisord group: %s' % group

def disable(config_name, check_exists=True):
    # check if symlink exists
    target = osp.join('servers', 'active', config_name+'.conf')
    if check_exists and not osp.islink(target):
        print 'ERROR: server not active: %s' % config_name
        return
    
    print 'Disabling "%s" config ...' % config_name

    # read the config file
    cp = configparser.ConfigParser()
    cp.readfp(open(target))

    # read all the groups
    sections = (g for g in cp.sections() if g.startswith('group'))
    groups = (g.split(':')[1] for g in sections)

    # remove the symlink
    print 'Deleting active symlink ...'
    try:
        os.unlink(target)
    except OSError:
        pass

    # stop and remove the groups from supervisor
    try:
        subprocess.check_call(['supervisorctl', 'reread'])
    except subprocess.CalledProcessError, e:
        print 'ERROR: could reread supervisord config'
        return
    for group in groups:
        print 'Stopping and removing %s ...' % group
        try:
            subprocess.check_call(['supervisorctl', 'stop', '%s:*' % group])
            subprocess.check_call(['supervisorctl', 'remove', group])
        except subprocess.CalledProcessError, e:
            print 'ERROR: could not stop / remove supervisord group: %s' % group

def restart(config_name):
    print 'Restarting "%s" config ...' % config_name
    target = osp.join('servers', 'active', config_name+'.conf')

    # read the config file
    cp = configparser.ConfigParser()
    cp.readfp(open(target))

    # read all the groups
    sections = (g for g in cp.sections() if g.startswith('group'))
    groups = (g.split(':')[1] for g in sections)
    
    # restart the groups from supervisor
    try:
        subprocess.check_call(['supervisorctl', 'reread'])
    except subprocess.CalledProcessError, e:
        print 'ERROR: could not reread supervisord config'
        return
    for group in groups:
        print 'Restarting %s ...' % group
        try:
            subprocess.check_call(['supervisorctl', 'restart', '%s:*' % group])
        except subprocess.CalledProcessError, e:
            print 'ERROR: could not restart supervisord group: %s' % group

def cleanup():
    # require root to set owner / group
    if os.getuid():
        raise ValueError('You must su or sudo to root to cleanup.')
    # confirm dir reset
    reset = _prompt('Are your sure you want to cleanup all backups?', 'n')
    if not reset:
        print 'Aborted'
        return
    print 'Deleting backup directories ...'
    for name in os.listdir('backups'):
        dname = osp.join('backups', name)
        if os.path.isdir(dname):
            shutil.rmtree(dname)

def run_from_args():
    usage = "usage: %prog [init|update|enable|disable|restart|remove|cleanup] [name]"
    parser = optparse.OptionParser(usage=usage)
    (options, args) = parser.parse_args()
    # never takes zero args
    if len(args) < 1:
        parser.print_usage()
        return

    # one arg params
    if args[0] == 'cleanup':
        cleanup()
        return
    
    # two arg params
    if len(args) < 2:
        parser.print_usage()
    elif args[0] == 'init':
        init(args[1])
    elif args[0] == 'update':
        update(args[1])
    elif args[0] == 'remove':
        remove(args[1])
    elif args[0] == 'enable':
        enable(args[1])
    elif args[0] == 'disable':
        disable(args[1])
    elif args[0] == 'restart':
        restart(args[1])
    else:
        parser.print_usage()

if __name__ == '__main__':
    run_from_args()