inter.diff

 aws/aws_manage_instances.py |   97 ++++++++++++++++++++++++++-----------------
 1 file changed, 59 insertions(+), 38 deletions(-)

diff --git a/aws/aws_manage_instances.py b/aws/aws_manage_instances.py
index 587e51c..fa0d2f8 100644
--- a/aws/aws_manage_instances.py
+++ b/aws/aws_manage_instances.py
@@ -1,19 +1,17 @@
 #!/usr/bin/env python
 
 import argparse
 import json
 import logging
 from time import gmtime, strftime
 from boto.ec2 import connect_to_region
 
-logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s")
 log = logging.getLogger(__name__)
-
 REGIONS = ['us-east-1', 'us-west-2']
 
 
 def start(i, dry_run):
     name = i.tags.get('Name', '')
     log.info("Starting %s..." % name)
     if dry_run:
         log.info("Dry run mode, skipping...")
@@ -34,90 +32,113 @@ def restart(i, dry_run):
     name = i.tags.get('Name', '')
     log.info("Restarting %s..." % name)
     if dry_run:
         log.info("Dry run mode, skipping...")
     else:
         i.reboot()
 
 
+def enable(i, dry_run):
+    name = i.tags.get('Name', '')
+    log.info("Enabling %s..." % name)
+    if dry_run:
+        log.info("Dry run mode, skipping...")
+    else:
+        # .add_tag overwrites existing tag
+        i.add_tag("moz-state", "ready")
+
+
+def disable(i, dry_run):
+    name = i.tags.get('Name', '')
+    moz_state = "disabled at %s" % strftime("%Y-%m-%d %H:%M:%S +0000",
+                                            gmtime())
+    log.info("Disabling %s, setting moz-state tag to '%s'..." % (name, moz_state))
+    if dry_run:
+        log.info("Dry run mode, skipping...")
+    else:
+        i.add_tag("moz-state", moz_state)
+
+
 def status(i):
     instance_id = i.id
     name = i.tags.get('Name', '')
     ip = i.private_ip_address
     state = i.state
-    print name, instance_id, ip, state
+    moz_state = i.tags.get('moz-state', '')
+    enabled = bool(moz_state == "ready")
+
+    print "Name:".rjust(8), name
+    print "ID:".rjust(8), instance_id
+    print "IP:".rjust(8), ip
+    print "Enabled:".rjust(8), enabled
+    print "State:".rjust(8), state
+    print "Tags:".rjust(8), ", ".join(["%s -> %s" % (k, v)
+                                       for k, v in i.tags.iteritems()])
+    print "=" * 72
 
 
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser()
-    parser.add_argument("-k", "--secrets", required=True,
-                        type=argparse.FileType('r'),
-                        help="file where secrets can be found")
+    parser.add_argument("-k", "--secrets", type=argparse.FileType('r'),
+                        help="optional file where secrets can be found")
+    parser.add_argument("-r", "--region", dest="regions", action="append",
+                        help="optional list of regions")
     parser.add_argument("action", choices=["stop", "start", "restart",
-                                           "status"],
+                                           "enable", "disable", "status"],
                         help="action to be performed")
-    state = parser.add_mutually_exclusive_group()
-    state.add_argument(
-        "--disable", action="store_true",
-        help="Set moz-state tag to a value which prevents automatic start"
-    )
-    state.add_argument(
-        "--enable", action="store_true",
-        help="Set moz-state tag to a value which enables automatic start"
-    )
     parser.add_argument("-n", "--dry-run", action="store_true",
                         help="Dry run mode")
     parser.add_argument("-q", "--quiet", action="store_true",
                         help="Supress logging messages")
     parser.add_argument("hosts", metavar="host", nargs="+",
                         help="hosts to be processed")
 
     args = parser.parse_args()
-    secrets = json.load(args.secrets)
+    if args.secrets:
+        secrets = json.load(args.secrets)
+    else:
+        secrets = None
 
+    logging.basicConfig(format="%(asctime)s - %(levelname)s - %(message)s")
     if not args.quiet:
         log.setLevel(logging.DEBUG)
     else:
         log.setLevel(logging.ERROR)
 
-    for region in REGIONS:
-        conn = connect_to_region(
-            region,
-            aws_access_key_id=secrets['aws_access_key_id'],
-            aws_secret_access_key=secrets['aws_secret_access_key']
-        )
+    if not args.regions:
+        args.regions = REGIONS
+    for region in args.regions:
+        if secrets:
+            conn = connect_to_region(
+                region,
+                aws_access_key_id=secrets['aws_access_key_id'],
+                aws_secret_access_key=secrets['aws_secret_access_key']
+            )
+        else:
+            conn = connect_to_region(region)
+
         res = conn.get_all_instances()
         instances = reduce(lambda a, b: a + b, [r.instances for r in res])
         for i in instances:
             name = i.tags.get('Name', '')
             instance_id = i.id
             if not i.private_ip_address:
                 # Terminated instances has no IP address assinged
                 log.warn("Skipping (terminated?) %s (%s)..." % (name,
                                                                 instance_id))
                 continue
             if name in args.hosts:
-                log.info("Found  %s (%s)..." % (name, instance_id))
-
-                moz_state = None
-                if args.enable:
-                    moz_state = "ready"
-                elif args.disable:
-                    moz_state = "disabled at %s" % strftime(
-                        "%Y-%m-%d %H:%M:%S +0000", gmtime())
-                if moz_state:
-                    log.info("Changing state to: %s" % moz_state)
-                    if args.dry_run:
-                        log.info("Dry run mode, not applying...")
-                    else:
-                        # .add_tag overwrites existing tag
-                        i.add_tag("moz-state", moz_state)
+                log.info("Found %s (%s)..." % (name, instance_id))
 
                 if args.action == "start":
                     start(i, args.dry_run)
                 elif args.action == "stop":
                     stop(i, args.dry_run)
                 elif args.action == "restart":
                     restart(i, args.dry_run)
+                elif args.action == "enable":
+                    enable(i, args.dry_run)
+                elif args.action == "disable":
+                    disable(i, args.dry_run)
                 elif args.action == "status":
                     status(i)