diff --git a/salt/modules/system.py b/salt/modules/system.py index 90a29f3c81..1421c3ddb6 100644 --- a/salt/modules/system.py +++ b/salt/modules/system.py @@ -7,6 +7,7 @@ from __future__ import absolute_import # Import python libs from datetime import datetime, timedelta, tzinfo import re +import os.path # Import salt libs import salt.utils @@ -428,3 +429,91 @@ class _FixedOffset(tzinfo): def dst(self, dt): # pylint: disable=W0613 return timedelta(0) + + +def _strip_quotes(str_q): + ''' + Helper function to strip off the ' or " off of a string + ''' + if str_q[0] == str_q[-1] and str_q.startswith(("'", '"')): + return str_q[1:-1] + return str_q + + +def get_computer_desc(): + ''' + Get PRETTY_HOSTNAME value stored in /etc/machine-info + If this file doesn't exist or the variable doesn't exist + return False. + + :return: Value of PRETTY_HOSTNAME if this does not exist False. + :rtype: str + + CLI Example: + + .. code-block:: bash + + salt '*' system.get_computer_desc + ''' + desc = None + hostname_cmd = salt.utils.which('hostnamectl') + if hostname_cmd: + desc = __salt__['cmd.run']('{0} status --pretty'.format(hostname_cmd)) + else: + pattern = re.compile(r'^\s*PRETTY_HOSTNAME=(.*)$') + try: + with salt.utils.fopen('/etc/machine-info', 'r') as mach_info: + for line in mach_info.readlines(): + match = pattern.match(line) + if match: + # get rid of whitespace then strip off quotes + desc = _strip_quotes(match.group(1).strip()).replace('\\"', '"') + # no break so we get the last occurance + except IOError: + return False + return desc + + +def set_computer_desc(desc): + ''' + Set PRETTY_HOSTNAME value stored in /etc/machine-info + This will create the file if it does not exist. If + it is unable to create or modify this file returns False. + + :param str desc: The computer description + :return: False on failure. True if successful. + + CLI Example: + + .. code-block:: bash + + salt '*' system.set_computer_desc "Michael's laptop" + ''' + hostname_cmd = salt.utils.which('hostnamectl') + if hostname_cmd: + result = __salt__['cmd.run']('{0} set-hostname --pretty {1}'.format(hostname_cmd, desc)) + return True if result == '' else False + + if not os.path.isfile('/etc/machine-info'): + f = salt.utils.fopen('/etc/machine-info', 'a') + f.close() + + is_pretty_hostname_found = False + pattern = re.compile(r'^\s*PRETTY_HOSTNAME=(.*)$') + new_line = 'PRETTY_HOSTNAME="{0}"'.format(desc.replace('"', '\\"')) + try: + with salt.utils.fopen('/etc/machine-info', 'r+') as mach_info: + lines = mach_info.readlines() + for i, line in enumerate(lines): + if pattern.match(line): + is_pretty_hostname_found = True + lines[i] = new_line + if not is_pretty_hostname_found: + lines.append(new_line) + # time to write our changes to the file + mach_info.seek(0, 0) + mach_info.write(''.join(lines)) + mach_info.write('\n') + return True + except IOError: + return False diff --git a/tests/integration/modules/system.py b/tests/integration/modules/system.py index 47215a99c5..b24f20b5fe 100644 --- a/tests/integration/modules/system.py +++ b/tests/integration/modules/system.py @@ -17,6 +17,7 @@ ensure_in_syspath('../../') # Import salt libs import integration import salt.utils +import salt.states.file @skipIf(not salt.utils.is_linux(), 'These tests can only be run on linux') @@ -29,6 +30,7 @@ class SystemModuleTest(integration.ModuleCase): def __init__(self, arg): super(self.__class__, self).__init__(arg) self._orig_time = None + self._machine_info = True def setUp(self): super(SystemModuleTest, self).setUp() @@ -44,6 +46,9 @@ class SystemModuleTest(integration.ModuleCase): if self._orig_time is not None: self._restore_time() self._orig_time = None + if self._machine_info is not True: + self._restore_machine_info() + self._machine_info = True def _save_time(self): self._orig_time = datetime.datetime.utcnow() @@ -64,6 +69,20 @@ class SystemModuleTest(integration.ModuleCase): ''' return abs(t1 - t2) < datetime.timedelta(seconds=seconds_diff) + def _save_machine_info(self): + if os.path.isfile('/etc/machine-info'): + with salt.utils.fopen('/etc/machine-info', 'r') as mach_info: + self._machine_info = mach_info.read() + else: + self._machine_info = False + + def _restore_machine_info(self): + if self._machine_info is not False: + with salt.utils.fopen('/etc/machine-info', 'w') as mach_info: + mach_info.write(self._machine_info) + else: + self.run_function('file.remove', ['/etc/machine-info']) + def test_get_system_date_time(self): ''' Test we are able to get the correct time @@ -205,6 +224,39 @@ class SystemModuleTest(integration.ModuleCase): self.assertTrue(result) self.assertTrue(self._same_times(time_now, cmp_time), msg=msg) + @skipIf(os.geteuid() != 0, 'you must be root to run this test') + def test_get_computer_desc(self): + ''' + Test getting the system hostname + ''' + res = self.run_function('system.get_computer_desc') + + hostname_cmd = salt.utils.which('hostnamectl') + if hostname_cmd: + desc = self.run_function('cmd.run', ["hostnamectl status --pretty"]) + self.assertEqual(res, desc) + else: + if not os.path.isfile('/etc/machine-info'): + self.assertFalse(res) + else: + with salt.utils.fopen('/etc/machine-info', 'r') as mach_info: + data = mach_info.read() + self.assertIn(res, data.decode('string_escape')) + + @destructiveTest + @skipIf(os.geteuid() != 0, 'you must be root to run this test') + def test_set_computer_desc(self): + ''' + Test setting the system hostname + ''' + self._save_machine_info() + desc = "test" + ret = self.run_function('system.set_computer_desc', [desc]) + computer_desc = self.run_function('system.get_computer_desc') + + self.assertTrue(ret) + self.assertEqual(computer_desc, desc) + if __name__ == '__main__': from integration import run_tests run_tests(SystemModuleTest)