From 5576dea216ce4b3ded085780bdfa32c8a82eeea5 Mon Sep 17 00:00:00 2001 From: Yoon Hong Date: Fri, 16 Sep 2016 20:02:43 -0400 Subject: [PATCH] Replace unnecessary shell out logics with python based ones (#232) --- azurelinuxagent/common/osutil/default.py | 15 ++---- azurelinuxagent/common/protocol/metadata.py | 8 ++- azurelinuxagent/common/utils/fileutil.py | 19 +++++-- azurelinuxagent/common/utils/textutil.py | 5 ++ azurelinuxagent/ga/exthandlers.py | 7 ++- azurelinuxagent/pa/deprovision/default.py | 4 +- azurelinuxagent/pa/provision/default.py | 2 +- tests/common/osutil/test_default.py | 22 ++++++++ tests/ga/test_extension.py | 2 +- tests/tools.py | 1 + tests/utils/test_file_util.py | 57 +++++++++++++++++++-- 11 files changed, 115 insertions(+), 27 deletions(-) diff --git a/azurelinuxagent/common/osutil/default.py b/azurelinuxagent/common/osutil/default.py index 2ffd8c2..73eaf23 100644 --- a/azurelinuxagent/common/osutil/default.py +++ b/azurelinuxagent/common/osutil/default.py @@ -16,6 +16,7 @@ # Requires Python 2.4+ and Openssl 1.0+ # +import multiprocessing import os import re import shutil @@ -761,19 +762,11 @@ class DefaultOSUtil(object): return base64.b64decode(data) def get_total_mem(self): - cmd = "grep MemTotal /proc/meminfo |awk '{print $2}'" - ret = shellutil.run_get_output(cmd) - if ret[0] == 0: - return int(ret[1])/1024 - else: - raise OSUtilError("Failed to get total memory: {0}".format(ret[1])) + # Get total memory in bytes and divide by 1024**2 to get the valu in MB. + return os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / (1024**2) def get_processor_cores(self): - ret = shellutil.run_get_output("grep 'processor.*:' /proc/cpuinfo |wc -l") - if ret[0] == 0: - return int(ret[1]) - else: - raise OSUtilError("Failed to get processor cores") + return multiprocessing.cpu_count() def set_admin_access_to_ip(self, dest_ip): #This allows root to access dest_ip diff --git a/azurelinuxagent/common/protocol/metadata.py b/azurelinuxagent/common/protocol/metadata.py index 448649a..071f1f7 100644 --- a/azurelinuxagent/common/protocol/metadata.py +++ b/azurelinuxagent/common/protocol/metadata.py @@ -16,6 +16,7 @@ # # Requires Python 2.4+ and Openssl 1.0+ +import base64 import json import os import shutil @@ -294,8 +295,11 @@ class Certificates(object): p7b_file = os.path.join(conf.get_lib_dir(), P7B_FILE_NAME) # Wrapping the certificate lines. - b64_cmd = "echo {0} | base64 -d > {1}" - shellutil.run(b64_cmd.format(data, p7b_file)) + # decode and save the result into p7b_file + fileStream = open(p7b_file, 'w') + fileStream.write(textutil.b64decode(data)) + fileStream.close() + ssl_cmd = "openssl pkcs7 -text -in {0} -inform der | grep -v '^-----' " ret, data = shellutil.run_get_output(ssl_cmd.format(p7b_file)) diff --git a/azurelinuxagent/common/utils/fileutil.py b/azurelinuxagent/common/utils/fileutil.py index 7ef4fef..b0b6fb7 100644 --- a/azurelinuxagent/common/utils/fileutil.py +++ b/azurelinuxagent/common/utils/fileutil.py @@ -21,11 +21,11 @@ File operation util functions """ +import glob import os import re import shutil import pwd -import tempfile import azurelinuxagent.common.logger as logger from azurelinuxagent.common.future import ustr import azurelinuxagent.common.utils.textutil as textutil @@ -111,9 +111,11 @@ def chmod(path, mode): os.chmod(path, mode) def rm_files(*args): - for path in args: - if os.path.isfile(path): - os.remove(path) + for paths in args: + #Find all possible file paths + for path in glob.glob(paths): + if os.path.isfile(path): + os.remove(path) def rm_dirs(*args): """ @@ -169,3 +171,12 @@ def findstr_in_file(file_path, pattern_str): return None +def get_all_files(root_path): + """ + Find all files under the given root path + """ + result = [] + for root, dirs, files in os.walk(root_path): + result.extend([os.path.join(root, file) for file in files]) + + return result diff --git a/azurelinuxagent/common/utils/textutil.py b/azurelinuxagent/common/utils/textutil.py index 6d460cb..db320a6 100644 --- a/azurelinuxagent/common/utils/textutil.py +++ b/azurelinuxagent/common/utils/textutil.py @@ -278,6 +278,11 @@ def b64encode(s): return base64.b64encode(bytes(s, 'utf-8')).decode('utf-8') return base64.b64encode(s) +def b64decode(s): + from azurelinuxagent.common.version import PY_VERSION_MAJOR + if PY_VERSION_MAJOR > 2: + return base64.b64decode(s).decode('utf-8') + return base64.b64decode(s) def safe_shlex_split(s): import shlex diff --git a/azurelinuxagent/ga/exthandlers.py b/azurelinuxagent/ga/exthandlers.py index 642fdf2..5f75dc5 100644 --- a/azurelinuxagent/ga/exthandlers.py +++ b/azurelinuxagent/ga/exthandlers.py @@ -21,6 +21,7 @@ import glob import json import os import shutil +import stat import subprocess import time import zipfile @@ -518,8 +519,10 @@ class ExtHandlerInstance(object): except IOError as e: raise ExtensionError(u"Failed to write and unzip plugin", e) - chmod = "find {0} -type f | xargs chmod u+x".format(self.get_base_dir()) - shellutil.run(chmod) + #Add user execute permission to all files under the base dir + for file in fileutil.get_all_files(self.get_base_dir()): + fileutil.chmod(file, os.stat(file).st_mode | stat.S_IXUSR) + self.report_event(message="Download succeeded") self.logger.info("Initialize extension directory") diff --git a/azurelinuxagent/pa/deprovision/default.py b/azurelinuxagent/pa/deprovision/default.py index 3a916e2..3d60c45 100644 --- a/azurelinuxagent/pa/deprovision/default.py +++ b/azurelinuxagent/pa/deprovision/default.py @@ -63,8 +63,8 @@ class DeprovisionHandler(object): def regen_ssh_host_key(self, warnings, actions): warnings.append("WARNING! All SSH host key pairs will be deleted.") - actions.append(DeprovisionAction(shellutil.run, - ['rm -f /etc/ssh/ssh_host_*key*'])) + actions.append(DeprovisionAction(fileutil.rm_files, + ['/etc/ssh/ssh_host_*key*'])) def stop_agent_service(self, warnings, actions): warnings.append("WARNING! The waagent service will be stopped.") diff --git a/azurelinuxagent/pa/provision/default.py b/azurelinuxagent/pa/provision/default.py index 39486c0..c7514c3 100644 --- a/azurelinuxagent/pa/provision/default.py +++ b/azurelinuxagent/pa/provision/default.py @@ -81,7 +81,7 @@ class ProvisionHandler(object): def reg_ssh_host_key(self): keypair_type = conf.get_ssh_host_keypair_type() if conf.get_regenerate_ssh_host_key(): - shellutil.run("rm -f /etc/ssh/ssh_host_*key*") + fileutil.rm_files("/etc/ssh/ssh_host_*key*") keygen_cmd = "ssh-keygen -N '' -t {0} -f /etc/ssh/ssh_host_{1}_key" shellutil.run(keygen_cmd.format(keypair_type, keypair_type)) thumbprint = self.get_ssh_host_key_thumbprint(keypair_type) diff --git a/tests/common/osutil/test_default.py b/tests/common/osutil/test_default.py index d9d00f6..d982b7e 100644 --- a/tests/common/osutil/test_default.py +++ b/tests/common/osutil/test_default.py @@ -142,5 +142,27 @@ class TestOSUtil(AgentTestCase): self.assertTrue(endpoint is not None) self.assertEqual(endpoint, "second") + def test_get_total_mem(self): + """ + Validate the returned value matches to the one retrieved by invoking shell command + """ + cmd = "grep MemTotal /proc/meminfo |awk '{print $2}'" + ret = shellutil.run_get_output(cmd) + if ret[0] == 0: + self.assertEqual(int(ret[1]) / 1024, get_osutil().get_total_mem()) + else: + self.fail("Cannot retrieve total memory using shell command.") + + def test_get_processor_cores(self): + """ + Validate the returned value matches to the one retrieved by invoking shell command + """ + cmd = "grep 'processor.*:' /proc/cpuinfo |wc -l" + ret = shellutil.run_get_output(cmd) + if ret[0] == 0: + self.assertEqual(int(ret[1]), get_osutil().get_processor_cores()) + else: + self.fail("Cannot retrieve number of process cores using shell command.") + if __name__ == '__main__': unittest.main() diff --git a/tests/ga/test_extension.py b/tests/ga/test_extension.py index 9e246b7..fe30ca1 100644 --- a/tests/ga/test_extension.py +++ b/tests/ga/test_extension.py @@ -162,7 +162,7 @@ class TestExtension(AgentTestCase): self.assertNotEquals(0, len(vm_status.vmAgent.extensionHandlers)) handler_status = vm_status.vmAgent.extensionHandlers[0] self.assertEquals(expected_status, handler_status.status) - self.assertEquals("OSTCExtensions.ExampleHandlerLinux", + self.assertEquals("OSTCExtensions.ExampleHandlerLinux", handler_status.name) self.assertEquals(version, handler_status.version) self.assertEquals(expected_ext_count, len(handler_status.extensions)) diff --git a/tests/tools.py b/tests/tools.py index 2732f76..8801a0c 100644 --- a/tests/tools.py +++ b/tests/tools.py @@ -56,6 +56,7 @@ class AgentTestCase(unittest.TestCase): def setUp(self): prefix = "{0}_".format(self.__class__.__name__) self.tmp_dir = tempfile.mkdtemp(prefix=prefix) + self.test_file = 'test_file' conf.get_autoupdate_enabled = Mock(return_value=True) conf.get_lib_dir = Mock(return_value=self.tmp_dir) ext_log_dir = os.path.join(self.tmp_dir, "azure") diff --git a/tests/utils/test_file_util.py b/tests/utils/test_file_util.py index 9a5479e..f16f409 100644 --- a/tests/utils/test_file_util.py +++ b/tests/utils/test_file_util.py @@ -24,8 +24,9 @@ from azurelinuxagent.common.future import ustr import azurelinuxagent.common.utils.fileutil as fileutil class TestFileOperations(AgentTestCase): + def test_read_write_file(self): - test_file=os.path.join(self.tmp_dir, 'test_file') + test_file=os.path.join(self.tmp_dir, self.test_file) content = ustr(uuid.uuid4()) fileutil.write_file(test_file, content) @@ -34,7 +35,7 @@ class TestFileOperations(AgentTestCase): os.remove(test_file) def test_rw_utf8_file(self): - test_file=os.path.join(self.tmp_dir, 'test_file') + test_file=os.path.join(self.tmp_dir, self.test_file) content = u"\u6211" fileutil.write_file(test_file, content, encoding="utf-8") @@ -43,14 +44,14 @@ class TestFileOperations(AgentTestCase): os.remove(test_file) def test_remove_bom(self): - test_file=os.path.join(self.tmp_dir, 'test_file') + test_file=os.path.join(self.tmp_dir, self.test_file) data = b'\xef\xbb\xbfhehe' fileutil.write_file(test_file, data, asbin=True) data = fileutil.read_file(test_file, remove_bom=True) self.assertNotEquals(0xbb, ord(data[0])) def test_append_file(self): - test_file=os.path.join(self.tmp_dir, 'test_file') + test_file=os.path.join(self.tmp_dir, self.test_file) content = ustr(uuid.uuid4()) fileutil.append_file(test_file, content) @@ -68,5 +69,53 @@ class TestFileOperations(AgentTestCase): filename = fileutil.base_name(filepath) self.assertEquals('abc', filename) + def test_remove_files(self): + import random + import string + import glob + random_word = lambda : ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(5)) + + #Create 10 test files + test_file = os.path.join(self.tmp_dir, self.test_file) + test_file2 = os.path.join(self.tmp_dir, 'another_file') + test_files = [test_file + random_word() for _ in range(5)] + \ + [test_file2 + random_word() for _ in range(5)] + for file in test_files: + open(file, 'a').close() + + #Remove files using fileutil.rm_files + test_file_pattern = test_file + '*' + test_file_pattern2 = test_file2 + '*' + fileutil.rm_files(test_file_pattern, test_file_pattern2) + + self.assertEqual(0, len(glob.glob(os.path.join(self.tmp_dir, test_file_pattern)))) + self.assertEqual(0, len(glob.glob(os.path.join(self.tmp_dir, test_file_pattern2)))) + + def test_get_all_files(self): + import random + import string + random_word = lambda: ''.join(random.choice(string.ascii_uppercase + string.digits) for _ in range(5)) + + # Create 10 test files at the root dir and 10 other in the sub dir + test_file = os.path.join(self.tmp_dir, self.test_file) + test_file2 = os.path.join(self.tmp_dir, 'another_file') + expected_files = [test_file + random_word() for _ in range(5)] + \ + [test_file2 + random_word() for _ in range(5)] + + test_subdir = os.path.join(self.tmp_dir, 'test_dir') + os.mkdir(test_subdir) + test_file_in_subdir = os.path.join(test_subdir, self.test_file) + test_file_in_subdir2 = os.path.join(test_subdir, 'another_file') + expected_files.extend([test_file_in_subdir + random_word() for _ in range(5)] + \ + [test_file_in_subdir2 + random_word() for _ in range(5)]) + + for file in expected_files: + open(file, 'a').close() + + # Get All files using fileutil.get_all_files + actual_files = fileutil.get_all_files(self.tmp_dir) + + self.assertEqual(set(expected_files), set(actual_files)) + if __name__ == '__main__': unittest.main()