image-tools: T5923: update system_console.py for new GRUB file structure

Add util function to set serial console speed in accordance with revised
GRUB file structure; in keeping with the intentions of the config_mode
script, adjust the GRUB var 'console_speed' to only modify ttyS0.

3 files changed, 64 insertions, 5 deletions

diff --git a/python/vyos/system/compat.py b/python/vyos/system/compat.py
index 319c3dabf..436da14e8 100644
--- a/python/vyos/system/compat.py
+++ b/python/vyos/system/compat.py
@@ -27,7 +27,7 @@ TMPL_GRUB_COMPAT: str = 'grub/grub_compat.j2'
 # define regexes and variables
 REGEX_VERSION = r'^menuentry "[^\n]*{\n[^}]*\s+linux /boot/(?P<version>\S+)/[^}]*}'
 REGEX_MENUENTRY = r'^menuentry "[^\n]*{\n[^}]*\s+linux /boot/(?P<version>\S+)/vmlinuz (?P<options>[^\n]+)\n[^}]*}'
-REGEX_CONSOLE = r'^.*console=(?P<console_type>[^\s\d]+)(?P<console_num>[\d]+).*$'
+REGEX_CONSOLE = r'^.*console=(?P<console_type>[^\s\d]+)(?P<console_num>[\d]+)(,(?P<console_speed>[\d]+))?.*$'
 REGEX_SANIT_CONSOLE = r'\ ?console=[^\s\d]+[\d]+(,\d+)?\ ?'
 REGEX_SANIT_INIT = r'\ ?init=\S*\ ?'
 REGEX_SANIT_QUIET = r'\ ?quiet\ ?'
@@ -131,6 +131,8 @@ def parse_entry(entry: tuple) -> dict:
     # find console type and number
     regex_filter = compile(REGEX_CONSOLE)
     entry_dict.update(regex_filter.match(entry[1]).groupdict())
+    speed = entry_dict.get('console_speed', None)
+    entry_dict['console_speed'] = speed if speed is not None else '115200'
     entry_dict['boot_opts'] = sanitize_boot_opts(entry[1])
 
     return entry_dict
@@ -271,9 +273,11 @@ def grub_cfg_fields(root_dir: str = '') -> dict:
         root_dir = disk.find_persistence()
 
     grub_cfg_main = f'{root_dir}/{grub.GRUB_CFG_MAIN}'
+    grub_vars = f'{root_dir}/{grub.CFG_VYOS_VARS}'
 
-    fields = {'default': 0, 'timeout': 5}
-    # 'default' and 'timeout' from legacy grub.cfg
+    fields = grub.vars_read(grub_vars)
+    # 'default' and 'timeout' from legacy grub.cfg resets 'default' to
+    # index, rather than uuid
     fields |= grub.vars_read(grub_cfg_main)
 
     fields['tools_version'] = SYSTEM_CFG_VER
diff --git a/python/vyos/system/grub.py b/python/vyos/system/grub.py
index a94729964..781962dd0 100644
--- a/python/vyos/system/grub.py
+++ b/python/vyos/system/grub.py
@@ -354,5 +354,18 @@ def set_console_type(console_type: str, root_dir: str = '') -> None:
     vars_current['console_type'] = str(console_type)
     vars_write(vars_file, vars_current)
 
-def set_raid(root_dir: str = '') -> None:
-    pass
+def set_console_speed(console_speed: str, root_dir: str = '') -> None:
+    """Write default console speed to GRUB configuration
+
+    Args:
+        console_speed (str): default console speed
+        root_dir (str, optional): an optional path to the root directory.
+        Defaults to empty.
+    """
+    if not root_dir:
+        root_dir = disk.find_persistence()
+
+    vars_file: str = f'{root_dir}/{CFG_VYOS_VARS}'
+    vars_current: dict[str, str] = vars_read(vars_file)
+    vars_current['console_speed'] = str(console_speed)
+    vars_write(vars_file, vars_current)
diff --git a/python/vyos/system/grub_util.py b/python/vyos/system/grub_util.py
new file mode 100644
index 000000000..9e79d41d4
--- /dev/null
+++ b/python/vyos/system/grub_util.py
@@ -0,0 +1,42 @@
+# Copyright 2024 VyOS maintainers and contributors <maintainers@vyos.io>
+#
+# This library is free software; you can redistribute it and/or
+# modify it under the terms of the GNU Lesser General Public
+# License as published by the Free Software Foundation; either
+# version 2.1 of the License, or (at your option) any later version.
+#
+# This library 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
+# Lesser General Public License for more details.
+#
+# You should have received a copy of the GNU Lesser General Public License
+# along with this library.  If not, see <http://www.gnu.org/licenses/>.
+
+from vyos.system import disk, grub, compat
+
+@compat.grub_cfg_update
+def set_console_speed(console_speed: str, root_dir: str = '') -> None:
+    """Write default console speed to GRUB configuration
+
+    Args:
+        console_speed (str): default console speed
+        root_dir (str, optional): an optional path to the root directory.
+        Defaults to empty.
+    """
+    if not root_dir:
+        root_dir = disk.find_persistence()
+
+    grub.set_console_speed(console_speed, root_dir)
+
+def update_console_speed(console_speed: str, root_dir: str = '') -> None:
+    """Update console_speed if different from current value"""
+
+    if not root_dir:
+        root_dir = disk.find_persistence()
+
+    vars_file: str = f'{root_dir}/{grub.CFG_VYOS_VARS}'
+    vars_current: dict[str, str] = grub.vars_read(vars_file)
+    console_speed_current = vars_current.get('console_speed', None)
+    if console_speed != console_speed_current:
+        set_console_speed(console_speed, root_dir)