1 files changed, 150 insertions, 2 deletions
diff --git a/python/vyos/component_version.py b/python/vyos/component_version.py
index 9662ebfcf..648d690b9 100644
--- a/python/vyos/component_version.py
+++ b/python/vyos/component_version.py
@@ -1,4 +1,4 @@
-# Copyright 2022 VyOS maintainers and contributors <maintainers@vyos.io>
+# Copyright 2022-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
@@ -35,14 +35,162 @@ VyOS 1.2:
 import os
 import re
 import sys
-import fileinput
+from dataclasses import dataclass
+from dataclasses import replace
+from typing import Optional
 
 from vyos.xml_ref import component_version
+from vyos.utils.file import write_file
 from vyos.version import get_version
 from vyos.defaults import directories
 
 DEFAULT_CONFIG_PATH = os.path.join(directories['config'], 'config.boot')
 
+REGEX_WARN_VYOS = r'(// Warning: Do not remove the following line.)'
+REGEX_WARN_VYATTA = r'(/\* Warning: Do not remove the following line. \*/)'
+REGEX_COMPONENT_VERSION_VYOS = r'// vyos-config-version:\s+"([\w@:-]+)"\s*'
+REGEX_COMPONENT_VERSION_VYATTA = r'/\* === vyatta-config-version:\s+"([\w@:-]+)"\s+=== \*/'
+REGEX_RELEASE_VERSION_VYOS = r'// Release version:\s+(\S*)\s*'
+REGEX_RELEASE_VERSION_VYATTA = r'/\* Release version:\s+(\S*)\s*\*/'
+
+CONFIG_FILE_VERSION = """\
+// Warning: Do not remove the following line.
+// vyos-config-version: "{}"
+// Release version: {}\n
+"""
+
+warn_filter_vyos = re.compile(REGEX_WARN_VYOS)
+warn_filter_vyatta = re.compile(REGEX_WARN_VYATTA)
+
+regex_filter = { 'vyos': dict(zip(['component', 'release'],
+                                  [re.compile(REGEX_COMPONENT_VERSION_VYOS),
+                                   re.compile(REGEX_RELEASE_VERSION_VYOS)])),
+                 'vyatta': dict(zip(['component', 'release'],
+                                    [re.compile(REGEX_COMPONENT_VERSION_VYATTA),
+                                     re.compile(REGEX_RELEASE_VERSION_VYATTA)])) }
+
+@dataclass
+class VersionInfo:
+    component: Optional[dict[str,int]] = None
+    release: str = get_version()
+    vintage: str = 'vyos'
+    config_body: Optional[str] = None
+    footer_lines: Optional[list[str]] = None
+
+    def component_is_none(self) -> bool:
+        return bool(self.component is None)
+
+    def config_body_is_none(self) -> bool:
+        return bool(self.config_body is None)
+
+    def update_footer(self):
+        f = CONFIG_FILE_VERSION.format(component_to_string(self.component),
+                                       self.release)
+        self.footer_lines = f.splitlines()
+
+    def update_syntax(self):
+        self.vintage = 'vyos'
+        self.update_footer()
+
+    def update_release(self, release: str):
+        self.release = release
+        self.update_footer()
+
+    def update_component(self, key: str, version: int):
+        if not isinstance(version, int):
+            raise ValueError('version must be int')
+        if self.component is None:
+            self.component = {}
+        self.component[key] = version
+        self.component = dict(sorted(self.component.items(), key=lambda x: x[0]))
+        self.update_footer()
+
+    def update_config_body(self, config_str: str):
+        self.config_body = config_str
+
+    def write_string(self) -> str:
+        config_body = '' if self.config_body is None else self.config_body
+        footer_lines = [] if self.footer_lines is None else self.footer_lines
+
+        return config_body + '\n' + '\n'.join(footer_lines) + '\n'
+
+    def write(self, config_file):
+        string = self.write_string()
+        try:
+            write_file(config_file, string, mode=0o660)
+        except Exception as e:
+            raise ValueError(e) from e
+
+def component_to_string(component: dict) -> str:
+    l = [f'{k}@{v}' for k, v in sorted(component.items(), key=lambda x: x[0])]
+    return ':'.join(l)
+
+def component_from_string(string: str) -> dict:
+    return {k: int(v) for k, v in re.findall(r'([\w,-]+)@(\d+)', string)}
+
+def version_info_from_file(config_file) -> VersionInfo:
+    version_info = VersionInfo()
+    try:
+        with open(config_file) as f:
+            config_str = f.read()
+    except OSError:
+        return None
+
+    if len(parts := warn_filter_vyos.split(config_str)) > 1:
+        vintage = 'vyos'
+    elif len(parts := warn_filter_vyatta.split(config_str)) > 1:
+        vintage = 'vyatta'
+    else:
+        version_info.config_body = parts[0] if parts else None
+        return version_info
+
+    version_info.vintage = vintage
+    version_info.config_body = parts[0]
+    version_lines = ''.join(parts[1:]).splitlines()
+    version_lines = [k for k in version_lines if k]
+    if len(version_lines) != 3:
+        raise ValueError(f'Malformed version strings: {version_lines}')
+
+    m = regex_filter[vintage]['component'].match(version_lines[1])
+    if not m:
+        raise ValueError(f'Malformed component string: {version_lines[1]}')
+    version_info.component = component_from_string(m.group(1))
+
+    m = regex_filter[vintage]['release'].match(version_lines[2])
+    if not m:
+        raise ValueError(f'Malformed component string: {version_lines[2]}')
+    version_info.release = m.group(1)
+
+    version_info.footer_lines = version_lines
+
+    return version_info
+
+def version_info_from_system() -> VersionInfo:
+    """
+    Return system component versions.
+    """
+    d = component_version()
+    sort_d = dict(sorted(d.items(), key=lambda x: x[0]))
+    version_info = VersionInfo(
+        component = sort_d,
+        release =  get_version(),
+        vintage = 'vyos'
+    )
+
+    return version_info
+
+def version_info_copy(v: VersionInfo) -> VersionInfo:
+    """
+    Make a copy of dataclass.
+    """
+    return replace(v)
+
+def version_info_prune_component(x: VersionInfo, y: VersionInfo) -> VersionInfo:
+    """
+    In place pruning of component keys of x not in y.
+    """
+    x.component = { k: v for k,v in x.component.items() if k in y.component }
+
 def from_string(string_line, vintage='vyos'):
     """
     Get component version dictionary from string.