Add actual code.

Author: ianmcorvidae

.gitignore

@@ -0,0 +1,3 @@
+config.yaml
+*.conf
+*~

config.yaml.sample

@@ -0,0 +1,16 @@
+_default:
+  enabled: 1
+  fail_timeout: 10s
+  max_fails: 10
+  port: 80
+fancycluster:
+  snap:
+    host: 192.168.0.2
+    weight: 18
+  crackle:
+    host: 192.168.0.3
+    weight: 22
+  pop:
+    enabled: 0
+    host: 192.168.0.4
+    weight: 19

fancycluster.conf.sample

@@ -0,0 +1,8 @@
+upstream fancycluster {
+    # snap
+    server 192.168.0.2:80 weight=18 max_fails=10 fail_timeout=10s;
+    # crackle
+    server 192.168.0.3:80 weight=22 max_fails=10 fail_timeout=10s;
+    # pop (disabled)
+    ## server 192.168.0.4:80 weight=19 max_fails=10 fail_timeout=10s;
+}

upstream_manager.py

@@ -0,0 +1,256 @@
+#!/usr/bin/env python2
+from yaml import load, dump
+try:
+    from yaml import CLoader as Loader, CDumper as Dumper
+except ImportError:
+    from yaml import Loader, Dumper
+from optparse import OptionParser
+from os.path import dirname, exists
+from os import remove
+
+###########
+# CLASSES #
+###########
+
+class Config:
+    """A class encapsulating the config.yaml config file"""
+    def __init__(self, filename=dirname(__file__) + '/config.yaml'):
+        self.filename = filename
+        self.config_raw = None
+        self.config = None
+        self.load()
+
+    def load(self):
+        if self.config_raw is None:
+            self.load_raw()
+        if self.config is None:
+            self.process()
+    def load_raw(self):
+        config_file = open(self.filename)
+        self.config_raw = load(config_file, Loader=Loader)
+        config_file.close()
+    def process(self):
+        config_raw = self.config_raw
+
+        config = dict([(cluster_key, dict([(server_key, config_raw.get('_default', {}).copy())
+                                           for server_key in config_raw[cluster_key].keys()
+                                           if server_key[0] != '_']))
+                       for cluster_key in config_raw.keys()
+                       if cluster_key[0] != '_'])
+        for cluster_key in config.keys():
+            for server_key in config[cluster_key].keys():
+                config[cluster_key][server_key].update(config_raw[cluster_key].get('_default', {}))
+                config[cluster_key][server_key].update(config_raw[cluster_key][server_key])
+
+            config[cluster_key]['_ip_hash'] = bool(config_raw[cluster_key].get('_ip_hash', False))
+            config[cluster_key]['_file'] = dirname(__file__) + '/' + config_raw[cluster_key].get('_file', cluster_key + '.conf')
+        self.config = config
+    def save(self):
+        config_file = open(self.filename, 'w')
+        config_file.write(dump(self.config_raw, default_flow_style=False, Dumper=Dumper))
+        config_file.close()
+
+    def cluster(self, cluster_name):
+        return Cluster(cluster_name, self.config[cluster_name])
+
+    def _set_prop(self, cluster, server, prop, value):
+        cluster_name = cluster.name
+        self.config[cluster_name][server][prop] = value
+
+        cluster_default = self.config_raw[cluster_name].get('_default', {}).get(prop, None)
+        global_default = self.config_raw.get('_default', {}).get(prop, None)
+        if cluster_default == value or (cluster_default is None and global_default == value):
+            if prop in self.config_raw[cluster_name][server]:
+                del self.config_raw[cluster_name][server][prop]
+        else:
+            self.config_raw[cluster_name][server][prop] = value
+
+    def enable(self, cluster, server):
+        self._set_prop(cluster, server, 'enabled', 1)
+    def disable(self, cluster, server):
+        self._set_prop(cluster, server, 'enabled', 0)
+
+    def backup(self, cluster, server):
+        self._set_prop(cluster, server, 'backup', 1)
+    def nonbackup(self, cluster, server):
+        self._set_prop(cluster, server, 'backup', 0)
+
+    def down(self, cluster, server):
+        self._set_prop(cluster, server, 'down', 1)
+    def up(self, cluster, server):
+        self._set_prop(cluster, server, 'down', 0)
+
+    def weight(self, cluster, server, new):
+        if new is None:
+            new = 1
+        self._set_prop(cluster, server, 'weight', new)
+    def max_fails(self, cluster, server, new):
+        if new is None:
+            new = 1
+        self._set_prop(cluster, server, 'max_fails', new)
+    def fail_timeout(self, cluster, server, new):
+        if new is None:
+            new = '10s'
+        self._set_prop(cluster, server, 'fail_timeout', new)
+
+class Cluster:
+    """A class encapsulating a specific upstream/cluster"""
+    def __init__(self, name, config):
+        self.name = name
+        self.ip_hash = config.get('_ip_hash')
+        self.filename = config.get('_file')
+        self.servers = [Server(name, config[name]) for name in config.keys() if name[0] != '_']
+    def save(self, **kwargs):
+        upstream_def = 'upstream %s {\n%%s}\n' % self.name
+        if self.ip_hash:
+            upstream_def = upstream_def % '    ip_hash;\n%s'
+        count = 1
+        for server in self.servers:
+            if kwargs.get('rotate', False) == count:
+                server.rotate = True
+            upstream_def = upstream_def % server.comment_line('    ', '\n%s')
+            upstream_def = upstream_def % server.upstream_line('    ', '\n%s')
+            if server.active():
+                count = count + 1
+        upstream_def = upstream_def % ''
+
+        config_file = open(self.filename, 'w')
+        config_file.write(upstream_def)
+        config_file.close()
+
+class Server:
+    """A class encapsulating a specific server within a cluster"""
+    def __init__(self, name, config):
+        self.name = name
+        if config.get('host', False):
+            self.host = config['host']
+            self.port = config['port']
+            self.upstream = "%s:%s" % (self.host, + self.port)
+        elif config.get('upstream', False):
+            self.host = self.port = None
+            self.upstream = config['upstream']
+        else:
+            raise Exception('Please provide either host/port or upstream for server %s' % self.name)
+        for prop in ['weight', 'max_fails', 'fail_timeout', 'down', 'backup', 'enabled']:
+            setattr(self, prop, config.get(prop, None))
+        self.rotate = False
+    def active(self):
+        return self.enabled and not self.down
+    def comment_line(self, prefix='', postfix=''):
+        ret = '# %s' % self.name
+        if not self.enabled:
+            ret = ret + ' (disabled)'
+        if self.down:
+            ret = ret + ' (down)'
+        if self.rotate:
+            ret = ret + ' (rotated out)'
+        return prefix + ret + postfix
+    def upstream_line(self, prefix='', postfix=''):
+        properties = []
+        if self.rotate:
+            properties.append('#')
+        elif not self.enabled:
+            properties.append('##')
+
+        properties.extend(['server', self.upstream])
+        for prop in ['weight', 'max_fails', 'fail_timeout']:
+            if getattr(self, prop) is not None:
+                properties.append("%s=%s" % (prop, getattr(self,prop)))
+        for prop in ['down', 'backup']:
+            if getattr(self, prop):
+                properties.append(prop)
+        return prefix + " ".join(properties) + ';' + postfix
+
+###########
+# ACTIONS #
+###########
+
+def rotate_action(config, cluster, *other_args):
+    statefile = dirname(__file__) + '/.rotate-' + cluster.name
+    if exists(statefile):
+        fh = open(statefile, 'r')
+        state = int(fh.read().strip()) + 1
+        fh.close()
+    else:
+        state = 1
+    valid = [server.host for server in cluster.servers if server.active()]
+    cluster.save(rotate=state)
+    if state > len(valid):
+        print "Done"
+        remove(statefile)
+    else:
+        print valid[state-1]
+        fh = open(statefile, 'w')
+        fh.write(str(state))
+        fh.close()
+
+def generate_action(config, cluster, *other_args):
+    cluster.save()
+    print "Saved"
+
+def disable_action(config, cluster, args, options, parser):
+    if len(args) < 1:
+        parser.error('Please provide a server to disable.')
+    to_disable = args[0]
+    if cluster.ip_hash:
+        config.down(cluster, to_disable)
+    else:
+        config.disable(cluster, to_disable)
+    config.save()
+    config.cluster(cluster.name).save()
+    print "Disabled " + to_disable
+
+def enable_action(config, cluster, args, options, parser):
+    if len(args) < 1:
+        parser.error('Please provide a server to enable.')
+    to_enable = args[0]
+    if cluster.ip_hash:
+        config.enable(cluster, to_enable)
+        config.up(cluster, to_enable)
+    else:
+        config.enable(cluster, to_enable)
+    config.save()
+    config.cluster(cluster.name).save()
+    print "Enabled " + to_enable
+
+def weight_action(config, cluster, args, options, parser):
+    if len(args) < 1:
+        parser.error('Please provide a server.')
+    if len(args) < 2:
+        parser.error('Please provide a new value.')
+    server = args[0]
+    value = args[1]
+    config.weight(cluster, server, value)
+    config.save()
+    config.cluster(cluster.name).save()
+    print "Changed %s weight to %s" % (server, value)
+
+actions = {'generate': generate_action,
+           'rotate': rotate_action,
+           'disable': disable_action,
+           'enable': enable_action,
+           'weight': weight_action}
+
+# Run the script!
+if __name__ == '__main__':
+    usage = """Usage: %%prog [options] cluster action [server ...]
+           'cluster' should be a named cluster in your config file
+           'action' should be one of: %s
+           'server' and past is required for some action types""" % ", ".join(actions.keys())
+    parser = OptionParser(usage=usage)
+    (options, args) = parser.parse_args()
+
+    if len(args) < 2:
+        parser.error('Please provide both a cluster and an action!')
+
+    config = Config()
+    if args[0] not in config.config.keys():
+        parser.error('Unknown cluster. Available clusters are: ' + ", ".join([k for k in config.keys() if k[0] != '_']))
+
+    cluster = config.cluster(args[0])
+    action = args[1]
+
+    if action not in actions.keys():
+        parser.error('Unknown action. Available actions are: ' + ", ".join(actions.keys()))
+
+    actions[action](config, cluster, args[2:], options, parser)