Commit 4cc93577 authored by alexAubin's avatar alexAubin

Rework the hook system to have more explicit I/O definitions and support python hooks

parent dd684d88
import json
import yaml
import logging
import shlex
import subprocess
from importlib import import_module
from pipes import quote
from django.conf import settings
logger = logging.getLogger(__name__)
"""
Hooks are defined in the settings with a syntax like :
settings.HOOKS = {"foo": {"CREATE_FILE": {"type": "bash",
"command": "touch {dir}/{file}"},
"CREATE_REMOTE_FILE": {"type": "bash",
"remote": "user@foo.com",
"command": "touch {dir}/{file}"},
"UPDATE_FILE": {"type": "bash",
"command": "tee {dir}/{file}",
"stdin": "{content}"},
"LOAD_FILE": {"type": "bash",
"command": "cat {dir}/{file}",
"parse_output": "yaml"},
"PROVISION": {"type": "bash",
"remote": "user@machine",
"command": "provision_foo.sh --name {name} --ipv4 {ipv4}",
"stdin": "{password}"}
},
"bar": {"PROVISION": {"type": "python",
"module": "bar",
"function": "provision"},
"GET_STATE": {"type": "python",
"module": "bar",
"function": "state"}
}
}
"""
class HookManager():
@staticmethod
......@@ -15,42 +49,10 @@ class HookManager():
Returns true or false if the given hook name has been defined in
settings (and is not empty string)
"""
return hasattr(settings, "HOOKS") and settings.HOOKS.get(module, {}).get(name, "")
@staticmethod
def get_and_format(module, name, **kwargs):
hook = settings.HOOKS[module][name]
# Here two possibilities :
# local command are just "foo --bar {bar}"
# remote command are ("user@machine", "foo --bar {bar}")
if isinstance(hook, basestring):
# N.B. : here we split() the string to have separate piece.
# So for example
# foo --bar "zblerg {bar}"
# becomes something like
# ["foo", "--bar", "zblerg {bar}"]
# and then format each piece independently.
#
# This is meant to avoid obvious injections. Though if you're using
# a hook, you should try to validate the format of the args before
# calling this (e.g. check that a variable domain provided by the
# user really looks like a domain)
command = shlex.split(hook)
return [e.format(**kwargs) for e in command]
elif isinstance(hook, tuple) and len(hook) == 2:
user_at_machine = hook[0]
remote_command = shlex.split(hook[1])
remote_command = [quote(e.format(**kwargs)) for e in remote_command]
remote_command = " ".join(remote_command)
return ["ssh", user_at_machine, remote_command]
else:
raise Exception("Invalid command format for hook %s for module %s, expected raw "
"string or 2-tuple (user@machine, command)" % (name, module))
return hasattr(settings, "HOOKS") and settings.HOOKS.get(module, {}).get(name, {})
@staticmethod
def run(module, name, stdin=None, **kwargs):
def run(module, name, **kwargs):
# The arg stdin can be provided and defines what will be fed on the
# standard input of the command
......@@ -63,17 +65,88 @@ class HookManager():
logger.debug("Running hook %s for module %s with args %s" % (name, module, kwargs))
command = HookManager.get_and_format(module, name, **kwargs)
hook = settings.HOOKS[module][name]
if hook["type"] == "bash":
(success, stdout, stderr) = HookManager._run_bash(hook, **kwargs)
elif hook["type"] == "python":
(success, stdout, stderr) = HookManager._run_python(hook, **kwargs)
return (success, stdout, stderr)
@staticmethod
def _format_hook_bash(hook, **kwargs):
command = hook["command"]
logger.debug("Hook command : %s" % command)
# N.B. : here we split() the string to have separate piece.
# So for example
# foo --bar "zblerg {bar}"
# becomes something like
# ["foo", "--bar", "zblerg {bar}"]
# and then format each piece independently.
#
# This is meant to avoid obvious injections. Though if you're using
# a hook, you should try to validate the format of the args before
# calling this (e.g. check that a variable domain provided by the
# user really looks like a domain)
to_run = [e.format(**kwargs) for e in shlex.split(command)]
if "remote" in hook:
user_at_machine = hook["remote"]
remote_command = [quote(e) for e in to_run]
remote_command = " ".join(remote_command)
to_run = ["ssh", user_at_machine, remote_command]
return to_run
@staticmethod
def _run_bash(hook, **kwargs):
# Format command with provided infos
to_run = HookManager._format_hook_bash(hook, **kwargs)
if "stdin" in hook:
stdin = hook["stdin"].format(**kwargs)
else:
stdin = None
logger.debug("Will run hook command : %s" % to_run)
p = subprocess.Popen(
command,
to_run,
stdin=subprocess.PIPE if stdin else None,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE)
stdout, stderr = p.communicate(stdin)
# TODO : for "GET_" hooks, parse output as yaml to turn it into a dict ?
success = p.returncode == 0
if "parse_output" in hook:
parse_format = hook["parse_output"]
if parse_format == "json":
try:
stdout = json.loads(stdout)
except Exception as e:
raise Exception("Failed to load stdout as json. Raw content:\n%s\nError:%s" % (stdout, e))
elif parse_format == "yaml":
try:
stdout = yaml.safe_load(stdout)
except Exception as e:
raise Exception("Failed to load stdout as yaml. Raw content:\n%s\nError:%s" % (stdout, e))
else:
raise Exception("Unsupported parse_format %s" % parse_format)
return (success, stdout, stderr)
@staticmethod
def _run_python(hook, **kwargs):
try:
module = import_module("custom_hooks." + hook["module"])
except Exception as e:
raise Exception("Could not load python hook %s : %s" % (hook["module"], e))
try:
output = getattr(module, hook["function"])(**kwargs)
return (True, output, None)
except Exception as e:
return (False, None, str(e))
......@@ -13,55 +13,119 @@ class SubscriptionTestCase(TestCase):
def setUp(self):
self.tmpdir = tempfile.mkdtemp()
settings.HOOKS = {"example": {
"TOUCH": "touch {tmpdir}/{file}",
"TEE": "tee {tmpdir}/{file}",
"SSH": ("user@example.com", "touch {dir}/{file}")}
}
settings.HOOKS = {"foo": {"CREATE_FILE": {"type": "bash",
"command": "touch {dir}/{file}"},
"CREATE_REMOTE_FILE": {"type": "bash",
"remote": "user@foo.com",
"command": "touch {dir}/{file}"},
"UPDATE_FILE": {"type": "bash",
"command": "tee {dir}/{file}",
"stdin": "{content}"},
"LOAD_FILE": {"type": "bash",
"command": "cat {dir}/{file}",
"parse_output": "yaml"},
"PROVISION": {"type": "bash",
"remote": "user@machine",
"command": "provision_foo.sh --name {name} --ipv4 {ipv4}",
"stdin": "{password}"}
},
"bar": {"PROVISION": {"type": "python",
"module": "bar",
"function": "provision"},
"GET_STATE": {"type": "python",
"module": "bar",
"function": "state"}
}
}
if not os.path.exists("../custom_hooks"):
os.makedirs("../custom_hooks/")
os.system("touch ../custom_hooks/__init__.py")
open("../custom_hooks/bar.py", "w").write("""
import os
def provision(**kwargs):
dir = kwargs["dir"]
file = kwargs["file"]
open("%s/%s" % (dir, file), "w").write("running")
def state(**kwargs):
dir = kwargs["dir"]
file = kwargs["file"]
path = "%s/%s" % (dir, file)
provisioned = os.path.exists(path)
state = open(path).read().strip() if provisioned else "N/A"
return {"provisioned": provisioned,
"state": state }
""")
def tearDown(self):
shutil.rmtree(self.tmpdir)
os.remove("../custom_hooks/bar.py")
def test_hook_is_defined(self):
self.assertTrue(HookManager.is_defined("example", "TOUCH"))
self.assertTrue(HookManager.is_defined("example", "TEE"))
self.assertTrue(HookManager.is_defined("example", "SSH"))
self.assertFalse(HookManager.is_defined("example", "WAT"))
self.assertTrue(HookManager.is_defined("foo", "CREATE_FILE"))
self.assertTrue(HookManager.is_defined("foo", "UPDATE_FILE"))
self.assertTrue(HookManager.is_defined("foo", "LOAD_FILE"))
self.assertFalse(HookManager.is_defined("foo", "WAT"))
def test_hook_format_hook_bash(self):
hook_create_file = settings.HOOKS["foo"]["CREATE_FILE"]
hook_update_file = settings.HOOKS["foo"]["UPDATE_FILE"]
hook_create_remote_file = settings.HOOKS["foo"]["CREATE_REMOTE_FILE"]
self.assertEqual(HookManager._format_hook_bash(hook_create_file, dir=self.tmpdir, file="foo"),
["touch", self.tmpdir + "/foo"])
self.assertEqual(HookManager._format_hook_bash(hook_update_file, dir=self.tmpdir, file="bar"),
["tee", self.tmpdir + "/bar"])
self.assertEqual(HookManager._format_hook_bash(hook_create_remote_file, dir=self.tmpdir, file="wat"),
["ssh", "user@foo.com", "touch %s/wat" % self.tmpdir])
def test_hook_get_and_format(self):
def test_hook_format_hook_bash_injection_attempt(self):
self.assertEqual(HookManager.get_and_format("example", "TOUCH", tmpdir=self.tmpdir, file="foo"),
["touch", self.tmpdir+"/foo"])
self.assertEqual(HookManager.get_and_format("example", "TEE", tmpdir=self.tmpdir, file="bar"),
["tee", self.tmpdir+"/bar"])
self.assertEqual(HookManager.get_and_format("example", "SSH", dir=self.tmpdir, file="wat"),
["ssh", "user@example.com", "touch %s/wat" % self.tmpdir])
hook_create_file = settings.HOOKS["foo"]["CREATE_FILE"]
hook_create_remote_file = settings.HOOKS["foo"]["CREATE_REMOTE_FILE"]
def test_hook_get_and_format_injection_attempt(self):
self.assertEqual(HookManager.get_and_format("example", "TOUCH", tmpdir="/foo", file="wat; rm -rf /"),
self.assertEqual(HookManager._format_hook_bash(hook_create_file, dir="/foo", file="wat; rm -rf /"),
["touch", "/foo/wat; rm -rf /"])
# ------------------
# ^ one single arg for touch
self.assertEqual(HookManager.get_and_format("example", "SSH", dir="/foo", file="wat; rm -rf /"),
["ssh", "user@example.com", "touch '/foo/wat; rm -rf /'"])
# ---------------
# ^ one single arg for touch
self.assertEqual(HookManager.get_and_format("example", "SSH", dir="/foo", file="wat\"; rm -rf /"),
["ssh", "user@example.com", "touch '/foo/wat\"; rm -rf /'"])
# ----------------------
# ^ one single arg for touch
self.assertEqual(HookManager.get_and_format("example", "SSH", dir="/foo", file="wat'; rm -rf /"),
["ssh", "user@example.com", "touch '/foo/wat'\"'\"'; rm -rf /'"])
# ---------------------------
# ^ one single arg for touch
def test_run(self):
self.assertEqual(HookManager.run("example", "TOUCH", tmpdir=self.tmpdir, file="foo"),
# ------------------ # noqa
# ^ one single arg for touch # noqa
self.assertEqual(HookManager._format_hook_bash(hook_create_remote_file, dir="/foo", file="wat; rm -rf /"),
["ssh", "user@foo.com", "touch '/foo/wat; rm -rf /'"])
# --------------- # noqa
# ^ one single arg for touch # noqa
self.assertEqual(HookManager._format_hook_bash(hook_create_remote_file, dir="/foo", file="wat\"; rm -rf /"),
["ssh", "user@foo.com", "touch '/foo/wat\"; rm -rf /'"])
# ---------------------- # noqa
# ^ one single arg for touch # noqa
self.assertEqual(HookManager._format_hook_bash(hook_create_remote_file, dir="/foo", file="wat'; rm -rf /"),
["ssh", "user@foo.com", "touch '/foo/wat'\"'\"'; rm -rf /'"])
# --------------------------- # noqa
# ^ one single arg for touch # noqa
def test_run_bash(self):
self.assertEqual(HookManager.run("foo", "CREATE_FILE", dir=self.tmpdir, file="foo"),
(True, "", ""))
self.assertTrue(os.path.exists(self.tmpdir+"/foo"))
self.assertTrue(os.path.exists(self.tmpdir + "/foo"))
self.assertEqual(HookManager.run("foo", "UPDATE_FILE", content="yolo: swag", dir=self.tmpdir, file="bar"),
(True, "yolo: swag", ""))
self.assertEqual(open(self.tmpdir + "/bar").read().strip(), "yolo: swag")
self.assertEqual(HookManager.run("foo", "LOAD_FILE", dir=self.tmpdir, file="bar"),
(True, {"yolo": "swag"}, ""))
def test_run_python(self):
self.assertEqual(HookManager.run("bar", "PROVISION", dir=self.tmpdir, file="foo"),
(True, None, None))
self.assertEqual(HookManager.run("bar", "GET_STATE", dir=self.tmpdir, file="foo"),
(True, {"provisioned": True, "state": "running"}, None))
self.assertEqual(HookManager.run("example", "TEE", stdin="pwet", tmpdir=self.tmpdir, file="bar"),
(True, "pwet", ""))
self.assertEqual(open(self.tmpdir+"/bar").read().strip(), "pwet")
self.assertEqual(HookManager.run("bar", "GET_STATE", dir=self.tmpdir, file="pwet"),
(True, {"provisioned": False, "state": "N/A"}, None))
......@@ -24,3 +24,4 @@ unidecode>=1.0,<1.1
django-debug-toolbar>=1.9.1,<1.10
# Higher requires Django>=1.9
django-extensions>=2.0,<2.1
pyyaml>=5.1,<5.2
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment