From 7da3a3c30552db6359abc5771a3260bceecdcb1a Mon Sep 17 00:00:00 2001 From: Steven Armstrong Date: Wed, 12 Oct 2011 14:30:10 +0200 Subject: [PATCH] remote code execution and tests Signed-off-by: Steven Armstrong --- lib/cdist/exec/__init__.py | 0 lib/cdist/exec/remote.py | 132 +++++++++++++++++++++++++++++ lib/cdist/test/exec/__init__.py | 0 lib/cdist/test/exec/test_remote.py | 111 ++++++++++++++++++++++++ 4 files changed, 243 insertions(+) create mode 100644 lib/cdist/exec/__init__.py create mode 100644 lib/cdist/exec/remote.py create mode 100644 lib/cdist/test/exec/__init__.py create mode 100644 lib/cdist/test/exec/test_remote.py diff --git a/lib/cdist/exec/__init__.py b/lib/cdist/exec/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/lib/cdist/exec/remote.py b/lib/cdist/exec/remote.py new file mode 100644 index 00000000..a5ad0845 --- /dev/null +++ b/lib/cdist/exec/remote.py @@ -0,0 +1,132 @@ +# -*- coding: utf-8 -*- +# +# 2011 Steven Armstrong (steven-cdist at armstrong.cc) +# 2011 Nico Schottelius (nico-cdist at schottelius.org) +# +# This file is part of cdist. +# +# cdist is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# cdist is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with cdist. If not, see . +# +# + +import io +import os +import sys +import subprocess +import logging + +import cdist + + +class RemoteScriptError(cdist.Error): + def __init__(self, script, command, script_content): + self.script = script + self.command = command + self.script_content = script_content + + def __str__(self): + return "Remote script execution failed: %s %s" % (self.script, self.command) + + +class Remote(object): + """Execute commands remotely. + + All interaction with the remote side should be done through this class. + Directly accessing the remote side from python code is a bug. + + """ + def __init__(self, target_host, remote_base_path, remote_exec, remote_copy): + self.target_host = target_host + self.base_path = remote_base_path + self._exec = remote_exec + self._copy = remote_copy + + self.conf_path = os.path.join(self.base_path, "conf") + self.object_path = os.path.join(self.base_path, "object") + + self.type_path = os.path.join(self.conf_path, "type") + self.global_explorer_path = os.path.join(self.conf_path, "explorer") + + self.log = logging.getLogger(self.target_host) + + def create_directories(self): + self.rmdir(self.base_path) + self.mkdir(self.base_path) + self.mkdir(self.conf_path) + + def rmdir(self, path): + """Remove directory on the remote side.""" + self.log.debug("Remote rmdir: %s", path) + self.run(["rm", "-rf", path]) + + def mkdir(self, path): + """Create directory on the remote side.""" + self.log.debug("Remote mkdir: %s", path) + self.run(["mkdir", "-p", path]) + + def transfer(self, source, destination): + """Transfer a file or directory to the remote side.""" + self.log.debug("Remote transfer: %s -> %s", source, destination) + self.rmdir(destination) + command = self._copy.split() + command.extend(["-r", source, self.target_host + ":" + destination]) + self.run_command(command) + + def run(self, command, env=None): + """Run the given command with the given environment on the remote side. + Return the output as a string. + + """ + # prefix given command with remote_exec + cmd = self._exec.split() + cmd.append(self.target_host) + cmd.extend(command) + return self.run_command(cmd, env=None) + + def run_command(self, command, env=None): + """Run the given command with the given environment. + Return the output as a string. + + """ + assert isinstance(command, (list, tuple)), "list or tuple argument expected, got: %s" % command + self.log.debug("Remote run: %s", command) + try: + return subprocess.check_output(command, env=env) + except subprocess.CalledProcessError: + raise cdist.Error("Command failed: " + " ".join(command)) + except OSError as error: + raise cdist.Error(" ".join(*args) + ": " + error.args[1]) + + def run_script(self, script, env=None): + """Run the given script with the given environment on the remote side. + Return the output as a string. + + """ + command = self._exec.split() + command.append(self.target_host) + command.extend(["/bin/sh", "-e"]) + command.append(script) + + self.log.debug("Remote run script: %s", command) + if env: + self.log.debug("Remote run script env: %s", env) + + try: + return subprocess.check_output(command, env=env) + except subprocess.CalledProcessError as error: + script_content = self.run(["cat", script]) + self.log.error("Code that raised the error:\n%s", script_content) + raise RemoteScriptError(script, command, script_content) + except EnvironmentError as error: + raise cdist.Error(" ".join(command) + ": " + error.args[1]) diff --git a/lib/cdist/test/exec/__init__.py b/lib/cdist/test/exec/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/lib/cdist/test/exec/test_remote.py b/lib/cdist/test/exec/test_remote.py new file mode 100644 index 00000000..ae5e7aed --- /dev/null +++ b/lib/cdist/test/exec/test_remote.py @@ -0,0 +1,111 @@ +# -*- coding: utf-8 -*- +# +# 2010-2011 Steven Armstrong (steven-cdist at armstrong.cc) +# +# This file is part of cdist. +# +# cdist is free software: you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by +# the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# cdist is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU General Public License for more details. +# +# You should have received a copy of the GNU General Public License +# along with cdist. If not, see . +# +# + +import unittest +import os +import tempfile +import getpass +import shutil +import string +import random + +import cdist +from cdist.exec import remote + + +class RemoteTestCase(unittest.TestCase): + + def mkdtemp(self, **kwargs): + return tempfile.mkdtemp(prefix='tmp.cdist.test.', **kwargs) + + def mkstemp(self, **kwargs): + return tempfile.mkstemp(prefix='tmp.cdist.test.', **kwargs) + + def setUp(self): + self.temp_dir = self.mkdtemp() + target_host = 'localhost' + remote_base_path = self.temp_dir + user = getpass.getuser() + remote_exec = "ssh -o User=%s -q" % user + remote_copy = "scp -o User=%s -q" % user + self.remote = remote.Remote(target_host, remote_base_path, remote_exec, remote_copy) + + def tearDown(self): + shutil.rmtree(self.temp_dir) + + def test_run_success(self): + self.remote.run(['/bin/true']) + + def test_run_fail(self): + self.assertRaises(cdist.Error, self.remote.run, ['/bin/false']) + + def test_run_script_success(self): + handle, script = self.mkstemp(dir=self.temp_dir) + fd = open(script, "w") + fd.writelines(["#!/bin/sh\n", "/bin/true"]) + fd.close() + self.remote.run_script(script) + + def test_run_script_fail(self): + handle, script = self.mkstemp(dir=self.temp_dir) + fd = open(script, "w") + fd.writelines(["#!/bin/sh\n", "/bin/false"]) + fd.close() + self.assertRaises(remote.RemoteScriptError, self.remote.run_script, script) + + def test_run_script_get_output(self): + handle, script = self.mkstemp(dir=self.temp_dir) + fd = open(script, "w") + fd.writelines(["#!/bin/sh\n", "echo foobar"]) + fd.close() + self.assertEqual(self.remote.run_script(script), b"foobar\n") + + def test_mkdir(self): + temp_dir = self.mkdtemp(dir=self.temp_dir) + os.rmdir(temp_dir) + self.remote.mkdir(temp_dir) + self.assertTrue(os.path.isdir(temp_dir)) + + def test_rmdir(self): + temp_dir = self.mkdtemp(dir=self.temp_dir) + self.remote.rmdir(temp_dir) + self.assertFalse(os.path.isdir(temp_dir)) + + def test_transfer_file(self): + handle, source = self.mkstemp(dir=self.temp_dir) + target = self.mkdtemp(dir=self.temp_dir) + self.remote.transfer(source, target) + self.assertTrue(os.path.isfile(target)) + + def test_transfer_dir(self): + source = self.mkdtemp(dir=self.temp_dir) + # put a file in the directory as payload + handle, source_file = self.mkstemp(dir=source) + source_file_name = os.path.split(source_file)[-1] + target = self.mkdtemp(dir=self.temp_dir) + self.remote.transfer(source, target) + # test if the payload file is in the target directory + self.assertTrue(os.path.isfile(os.path.join(target, source_file_name))) + + def test_create_directories(self): + self.remote.create_directories() + self.assertTrue(os.path.isdir(self.remote.base_path)) + self.assertTrue(os.path.isdir(self.remote.conf_path))