power_state: support 'condition' argument

  
  if 'condition' is provided to config in power_state, then
  consult it before powering off.
  
  This allows the user to shut down only if a condition is met, and
  leave the system in a debuggable state otherwise.
  
  An example is as simple as:
   power_state:
     mode: poweroff
     condition: ['sh', '-c', '[ -f /disable-poweroff ]']

4 files changed, 105 insertions, 10 deletions

diff --git a/ChangeLog b/ChangeLog
index 6fb70696..bbb7e990 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -61,6 +61,7 @@
  - status_wrapper in main: fix use of print_exc when handling exception
  - reporting: add reporting module for web hook or logging of events.
  - NoCloud: fix consumption of vendordata (LP: #1493453)
+ - power_state_change: support 'condition' to disable or enable poweroff
 0.7.6:
  - open 0.7.6
  - Enable vendordata on CloudSigma datasource (LP: #1303986)
diff --git a/cloudinit/config/cc_power_state_change.py b/cloudinit/config/cc_power_state_change.py
index 09d37371..7d9567e3 100644
--- a/cloudinit/config/cc_power_state_change.py
+++ b/cloudinit/config/cc_power_state_change.py
@@ -22,6 +22,7 @@ from cloudinit import util
 import errno
 import os
 import re
+import six
 import subprocess
 import time
 
@@ -48,10 +49,40 @@ def givecmdline(pid):
         return None
 
 
+def check_condition(cond, log=None):
+    if isinstance(cond, bool):
+        if log:
+            log.debug("Static Condition: %s" % cond)
+        return cond
+
+    pre = "check_condition command (%s): " % cond
+    try:
+        proc = subprocess.Popen(cond, shell=not isinstance(cond, list))
+        proc.communicate()
+        ret = proc.returncode
+        if ret == 0:
+            if log:
+                log.debug(pre + "exited 0. condition met.")
+            return True
+        elif ret == 1:
+            if log:
+                log.debug(pre + "exited 1. condition not met.")
+            return False
+        else:
+            if log:
+                log.warn(pre + "unexpected exit %s. " % ret +
+                         "do not apply change.")
+            return False
+    except Exception as e:
+        if log:
+            log.warn(pre + "Unexpected error: %s" % e)
+        return False
+
+
 def handle(_name, cfg, _cloud, log, _args):
 
     try:
-        (args, timeout) = load_power_state(cfg)
+        (args, timeout, condition) = load_power_state(cfg)
         if args is None:
             log.debug("no power_state provided. doing nothing")
             return
@@ -59,6 +90,10 @@ def handle(_name, cfg, _cloud, log, _args):
         log.warn("%s Not performing power state change!" % str(e))
         return
 
+    if condition is False:
+        log.debug("Condition was false. Will not perform state change.")
+        return
+
     mypid = os.getpid()
 
     cmdline = givecmdline(mypid)
@@ -70,8 +105,8 @@ def handle(_name, cfg, _cloud, log, _args):
 
     log.debug("After pid %s ends, will execute: %s" % (mypid, ' '.join(args)))
 
-    util.fork_cb(run_after_pid_gone, mypid, cmdline, timeout, log, execmd,
-                 [args, devnull_fp])
+    util.fork_cb(run_after_pid_gone, mypid, cmdline, timeout, log, 
+                 condition, execmd, [args, devnull_fp])
 
 
 def load_power_state(cfg):
@@ -80,7 +115,7 @@ def load_power_state(cfg):
     pstate = cfg.get('power_state')
 
     if pstate is None:
-        return (None, None)
+        return (None, None, None)
 
     if not isinstance(pstate, dict):
         raise TypeError("power_state is not a dict.")
@@ -115,7 +150,10 @@ def load_power_state(cfg):
         raise ValueError("failed to convert timeout '%s' to float." %
                          pstate['timeout'])
 
-    return (args, timeout)
+    condition = pstate.get("condition", True)
+    if not isinstance(condition, six.string_types + (list, bool)):
+        raise TypeError("condition type %s invalid. must be list, bool, str")
+    return (args, timeout, condition)
 
 
 def doexit(sysexit):
@@ -133,7 +171,7 @@ def execmd(exe_args, output=None, data_in=None):
     doexit(ret)
 
 
-def run_after_pid_gone(pid, pidcmdline, timeout, log, func, args):
+def run_after_pid_gone(pid, pidcmdline, timeout, log, condition, func, args):
     # wait until pid, with /proc/pid/cmdline contents of pidcmdline
     # is no longer alive.  After it is gone, or timeout has passed
     # execute func(args)
@@ -175,4 +213,11 @@ def run_after_pid_gone(pid, pidcmdline, timeout, log, func, args):
 
     if log:
         log.debug(msg)
+
+    try:
+        if not check_condition(condition, log):
+            return
+    except Exception as e:
+        fatal("Unexpected Exception when checking condition: %s" % e)
+
     func(*args)
diff --git a/doc/examples/cloud-config-power-state.txt b/doc/examples/cloud-config-power-state.txt
index 8df14366..b470153d 100644
--- a/doc/examples/cloud-config-power-state.txt
+++ b/doc/examples/cloud-config-power-state.txt
@@ -23,9 +23,18 @@
 # message: provided as the message argument to 'shutdown'. default is none.
 # timeout: the amount of time to give the cloud-init process to finish
 #          before executing shutdown.
+# condition: apply state change only if condition is met.
+#            May be boolean True (always met), or False (never met),
+#            or a command string or list to be executed.
+#            command's exit code indicates:
+#               0: condition met
+#               1: condition not met
+#            other exit codes will result in 'not met', but are reserved
+#            for future use.
 #
 power_state:
  delay: "+30"
  mode: poweroff
  message: Bye Bye
  timeout: 30
+ condition: True
diff --git a/tests/unittests/test_handler/test_handler_power_state.py b/tests/unittests/test_handler/test_handler_power_state.py
index 2f86b8f8..5687b10d 100644
--- a/tests/unittests/test_handler/test_handler_power_state.py
+++ b/tests/unittests/test_handler/test_handler_power_state.py
@@ -1,6 +1,9 @@
+import sys
+
 from cloudinit.config import cc_power_state_change as psc
 
 from .. import helpers as t_help
+from ..helpers import mock
 
 
 class TestLoadPowerState(t_help.TestCase):
@@ -9,12 +12,12 @@ class TestLoadPowerState(t_help.TestCase):
 
     def test_no_config(self):
         # completely empty config should mean do nothing
-        (cmd, _timeout) = psc.load_power_state({})
+        (cmd, _timeout, _condition) = psc.load_power_state({})
         self.assertEqual(cmd, None)
 
     def test_irrelevant_config(self):
         # no power_state field in config should return None for cmd
-        (cmd, _timeout) = psc.load_power_state({'foo': 'bar'})
+        (cmd, _timeout, _condition) = psc.load_power_state({'foo': 'bar'})
         self.assertEqual(cmd, None)
 
     def test_invalid_mode(self):
@@ -53,23 +56,60 @@ class TestLoadPowerState(t_help.TestCase):
     def test_no_message(self):
         # if message is not present, then no argument should be passed for it
         cfg = {'power_state': {'mode': 'poweroff'}}
-        (cmd, _timeout) = psc.load_power_state(cfg)
+        (cmd, _timeout, _condition) = psc.load_power_state(cfg)
         self.assertNotIn("", cmd)
         check_lps_ret(psc.load_power_state(cfg))
         self.assertTrue(len(cmd) == 3)
 
+    def test_condition_null_raises(self):
+        cfg = {'power_state': {'mode': 'poweroff', 'condition': None}}
+        self.assertRaises(TypeError, psc.load_power_state, cfg)
+
+    def test_condition_default_is_true(self):
+        cfg = {'power_state': {'mode': 'poweroff'}}
+        _cmd, _timeout, cond = psc.load_power_state(cfg)
+        self.assertEqual(cond, True)
+
+
+class TestCheckCondition(t_help.TestCase):
+    def cmd_with_exit(self, rc):
+        return([sys.executable, '-c', 'import sys; sys.exit(%s)' % rc])
+        
+    def test_true_is_true(self):
+        self.assertEqual(psc.check_condition(True), True)
+
+    def test_false_is_false(self):
+        self.assertEqual(psc.check_condition(False), False)
+
+    def test_cmd_exit_zero_true(self):
+        self.assertEqual(psc.check_condition(self.cmd_with_exit(0)), True)
+
+    def test_cmd_exit_one_false(self):
+        self.assertEqual(psc.check_condition(self.cmd_with_exit(1)), False)
+
+    def test_cmd_exit_nonzero_warns(self):
+        mocklog = mock.Mock()
+        self.assertEqual(
+            psc.check_condition(self.cmd_with_exit(2), mocklog), False)
+        self.assertEqual(mocklog.warn.call_count, 1)
+
+
 
 def check_lps_ret(psc_return, mode=None):
-    if len(psc_return) != 2:
+    if len(psc_return) != 3:
         raise TypeError("length returned = %d" % len(psc_return))
 
     errs = []
     cmd = psc_return[0]
     timeout = psc_return[1]
+    condition = psc_return[2]
 
     if 'shutdown' not in psc_return[0][0]:
         errs.append("string 'shutdown' not in cmd")
 
+    if 'condition' is None:
+        errs.append("condition was not returned")
+
     if mode is not None:
         opt = {'halt': '-H', 'poweroff': '-P', 'reboot': '-r'}[mode]
         if opt not in psc_return[0]: