Merge pull request #35482 from raazvvann/dev/rheghedu/system-comments

Add linux support for {get,set}_computer_desc CONTINUED

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

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)